Автор | BBuf
Оригинальный текст был впервые опубликован на GiantPandaCV.
1
предисловие
Недавно, с помощью моего коллеги shenghang, я сделал небольшую разработку, связанную с OneFlow IR, и у меня есть некоторые новые ощущения от части выполнения MLIR, поэтому я пытаюсь поделиться ею. Я провел много времени, изучая всю архитектуру OneFlow IR (см. мою серию учебных пособий по игрушкам), но я скептически отнесся к реализации JIT-части OneFlow IR. В последнее время OneFlow основан на Job (рабочая функция OneFlow, которую можно понимать как расчетный график без учета оборудования) в части реализации проекта MLIR и реорганизован, и под руководством Shenghang я понял весь процесс.
Итак, в этом документе я расскажу, как объединяются OneFlow и MLIR, как добавить проход на уровне графа в OneFlow IR, как операция OneFlow автоматически становится операцией MLIR и почему OneFlow IR может использовать MLIR для ускорения вычислений и т. д. Я мало что знаю о MLIR.Я начал с ним связываться 2 месяца назад.Если есть какие-то ошибки, прошу покритиковать и исправить.
Эта статья иGitHub.com/oneflow-Inc… & GitHub.com/B buf/TVs_beautiful…Если вам интересно, можете поставить звездочку и обратить внимание.
Op и Operation, упомянутые в этой статье, — это одно и то же, и строгого различия между ними нет.
2
Как OneFlow интегрируется с MLIR?
Внедрение MLIR в OneFlow в качестве IR OneFlow имеет много преимуществ: оно может не только заменить определение операции, которое необходимо написать вручную с помощью C++ в OneFlow, чтобы уменьшить сложность разработки, но и уменьшить некоторые накладные расходы, связанные с контейнером, в определении операции. . Кроме того, мы также можем ускорить расчет вычислительного графа с помощью инфраструктуры, поддерживаемой MLIR (т. е. нескольких диалектов).
Вычислительный граф здесь может быть либо вычислительным графом Игера, либо вычислительным графом Ленивца. Поскольку работа по использованию MLIR для ускорения на основе вычислительного графа Eager (то есть oneflow.jit.xxx) официально не открывалась, я все же использую Lazy вычислительный граф (Job) в качестве примера для объяснения процесса объединения OneFlow. и МЛИР.
Во-первых, нам нужно скомпилировать OneFlow с включенным MLIR, команда компиляции выглядит следующим образом:
git clone git@github.com:Oneflow-Inc/oneflow.git
cd oneflow && mkdir build && cd build
cmake-C ../cmake/caches/cn/fast/mlir-cuda-75.cmake -DBUILD_TESTING=ON .. && ninja
Затем вы можете написать пример для тестирования:
os.environ["ONEFLOW_MLIR_ENABLE_ROUND_TRIP"] = '1'
os.environ["ONEFLOW_MLIR_ENABLE_CODEGEN_FUSERS"] = '1'
@flow.unittest.skip_unless_1n1d()
class TestFuseBiasAddGeLUCPUMLIR(oneflow.unittest.TestCase):
def test_fused_bias_add_gelu_graph(test_case):
data = np.random.randn(1, 2, 3)
bias_data = np.random.randn(2)
x = flow.tensor(data, dtype=flow.float32)
bias = flow.tensor(bias_data, dtype=flow.float32)
y_eager = flow.gelu(flow._C.bias_add(x, bias, axis=1))
class FuseBiasAddGeLUGraph(flow.nn.Graph):
def __init__(self):
super().__init__()
def build(self, x):
return flow.gelu(flow._C.bias_add(x, bias, axis=1))
bias_add_gelu = FuseBiasAddGeLUGraph()
y_lazy = bias_add_gelu(x)
test_case.assertTrue(np.array_equal(y_eager.numpy(), y_lazy.numpy()))
После запуска этого примера в текущем рабочем каталоге будет создан файл журнала, а также будет папка ir_pass, в которой записан график расчета (.prototxt) до и после оптимизации OneFlow MLIR и выражение MLIR (.mlir), иФайлы .mlir.dot можно открывать с помощью graphviz для визуализации вычислительных графиков выражений MLIR.
Следует отметить, что если OneFlow выполняет обучающую задачу, эта папка журнала содержит не только графики прямого расчета и выражения MLIR, но также создает графики обратного расчета и выражения MLIR. Таким образом, MLIR может играть роль в процессе работы всей нейронной сети, что является важным моментом, отличающимся от системы прямого рассуждения, то есть обучение также может быть ускорено.
В oneflow/api/python/ir.cpp есть следующие две строки кода:
REGISTER_JOB_PASS("IRRoundTripBeforeAD", IRRoundTrip<kBeforeAD>);
REGISTER_JOB_PASS("IRRoundTrip", IRRoundTrip<kAfterAD>);
RoundTrip означает круговой путь. BeforeAD можно понимать как до обратного, а kAfterAD можно понимать как после обратного. Здесь связь между графом расчета OneFlow и MLIR устанавливается путем регистрации процесса взаимного преобразования между OneFlow Job и MLIR как Прохождение задания OneFlow. При выполнении скрипта OneFlow, если вы хотите, чтобы MLIR воздействовал на график вычислений OneFlow, вы можете включить переменную среды ONEFLOW_MLIR_ENABLE_ROUND_TRIP=1.
Далее, установление связи между графом расчета OneFlow и MLIR эквивалентно однозначному преобразованию между Операцией в графе расчета OneFlow и Операцией в MLIR. Операция MLIR определена в разделе «Диалекты» на всех уровнях.Согласно принципу общего доступа MLIR, мы внедрили диалект OneFlow и реализовали взаимно однозначное сопоставление операции OneFlow с операцией в рамках диалекта OneFlow на диалекте OneFlow.
Как определить диалект и операцию OneFlow, здесь обсуждаться не будет.Вы можете обратиться к разделу «Диалекты и ODS» официального документа MLIR (Маори. Ли Луму.org/docs/op def i…Определение операции Я ранее резюмировал документ в сочетании с определением операции диалекта OneFlow (GitHub.com/B buf/TVs_beautiful…середина) .
В дополнение к определениям диалекта и операции необходимо определить некоторые другие вещи, например, сопоставление типа данных OneFlow с типом данных MLIR определено в файле oneflow/ir/include/OneFlow/OneFlowEnums.td и некоторые общие внешние интерфейсы OneFlow Dialect Operation определены в файле oneflow /ir/include/OneFlow/OneFlowEnums.td. Здесь мы возьмем операцию изменения формы в качестве примера, чтобы кратко проиллюстрировать компоненты этой операции:
def OneFlow_ReshapeOp : OneFlow_BaseOp<"reshape", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
let input = (ins
AnyType:$in
);
let output = (outs
AnyType:$out
);
let attrs = (ins
AnyI64ElementsAttr:$shape
);
}
Имя OneFlow_ReshapeOp перед символом подчеркивания — это имя диалекта, за которым следует имя операции в диалекте. Затем эта операция наследует базовый класс OneFlow_BaseOp и объявляет ограничения и внешние интерфейсы, а затем определяет ввод, вывод и свойства операции.
Можно обнаружить, что определение диалектной операции OneFlow точно такое же, как и определение пользовательской операции OneFlow, что обеспечивает законность взаимной передачи между OneFlow и MLIR. Операция OneFlow Reshape определяется следующим образом:
REGISTER_USER_OP("reshape")
.Input("in")
.Output("out")
.Attr<Shape>("shape")
...
Взаимное преобразование между OneFlow Job и MLIR реализовано в oneflow/ir/oneflow-translate.Главное, что нужно сделать, это пройтись по OpGraph Job, обработать узлы и ребра соответственно и, наконец, преобразовать их в выражение MLIR. в то же время, после завершения расчета, он может быть основан на MLIR.Expression переопределяет Job. Общая логика здесь сложная, потому что необходимо иметь дело с преобразованием различных типов операций и ребер в OneFlow Job OpGraph, Я не буду продолжать подробно объяснять это здесь, потому что это не тот вопрос, который я хочу обсудить. в этой статье.Если вам интересно, вы можете прочитать код напрямую.
3
Как работает OneFlow IR?
В определении эксплуатации приведен пример изменить. Это легко просматривать OneFlow / IR / включать / OneFlow / OneFlowops.td, чтобы найти, что oneflow_mlirjitop также определен здесь. Этот пользовательский op используется для выполнения выражений MLIR. Ядро, которое реализует CPU и GPU (исходный код в OneFlow / IR / OneFlow-Extension / Paillow.cpp) используется для загрузки двигателя JIT-исполнения, предоставленного MLIR для запуска Final LVM IR. Итак, как произошла LLVM IR? Это получается по спуску экспрессии MLIR OneFlow MLIR по шагу. Конкретный процесс по убыванию выглядит следующим образом:
void AddLowerToLinalgMemRefPasses(PassManager& pm) {
pm.addPass(createLowerOneFlowToTosaPass()); // lower-oneflow-to-tosa
pm.addPass(createCSEPass()); // cse
pm.addNestedPass<FuncOp>(tosa::createTosaToLinalg()); // tosa-to-linalg-on-tensors
auto p = createLinalgElementwiseOpFusionPass();
assert(p->initializeOptions("allow-folding-unit-dim-reshapes=true").succeeded());
pm.addNestedPass<FuncOp>(std::move(p)); // linalg-fuse-elementwise-ops
pm.addNestedPass<FuncOp>(createLinalgBufferizePass()); // linalg-bufferize
pm.addNestedPass<FuncOp>(createTensorBufferizePass()); // tensor-bufferize
pm.addPass(createTensorConstantBufferizePass()); // tensor-constant-bufferize
pm.addPass(createFuncBufferizePass()); // func-bufferize
pm.addPass(createBufferResultsToOutParamsPass()); // buffer-results-to-out-params
pm.addPass(createCanonicalizerPass()); // canonicalize
pm.addNestedPass<FuncOp>(createFinalizingBufferizePass()); // finalizing-bufferize
}
LogicalResult LowerModuleToLLVM(mlir::MLIRContext* context, ModuleOp module) {
mlir::PassManager pm(context);
AddLowerToLinalgMemRefPasses(pm);
pm.addNestedPass<FuncOp>(createConvertLinalgToLoopsPass()); // convert-linalg-to-loops
pm.addNestedPass<FuncOp>(createLowerToCFGPass()); // convert-scf-to-std
pm.addPass(createConvertLinalgToLLVMPass()); // convert-linalg-to-llvm
pm.addPass(createMemRefToLLVMPass()); // convert-memref-to-llvm
pm.addPass(createLowerToLLVMPass()); // convert-std-to-llvm
pm.addPass(createReconcileUnrealizedCastsPass());
return pm.run(module);
}
Можно увидеть, как диалект OneFlow спускается сначала к диалекту Tosa, затем к диалекту Linalg, затем к диалекту Loop и вплоть до финального LLVM IR. В процессе пошагового спуска мы можем пользоваться возможностями оптимизации, предоставляемыми преобразованиями вложенных циклов, такими как Linalg Dialect, для повышения производительности конечного IR.
Процесс понижения здесь запускается, когда OneFlow вызывает ядро MlirJitOp (oneflow/ir/oneflow-extension/extension.cpp), и вызов также добавляется в процесс оптимизации как проход MLIR. Реализация процесса вызова JIT Pass может быть упрощена следующим образом:
class OutlineJitFunctionPass : public OutlineJitFunctionPassBase<OutlineJitFunctionPass> {
void runOnOperation() override {
Operation* op = getOperation();
RewritePatternSet patterns(op->getContext());
oneflow::populateFuserPasses(patterns);
(void)applyPatternsAndFoldGreedily(op, std::move(patterns));
}
};
std::unique_ptr<Pass> createOutlineJitFunctionPass() {
return std::make_unique<OutlineJitFunctionPass>();
}
LogicalResult ApplyRoundTripPatterns(RoundTripOneFlowJobWrapperInterface& job_wrapper,
MLIRContext* context, OwningModuleRef& module) {
mlir::PassManager pm(context);
pm.addNestedPass<mlir::FuncOp>(::mlir::createCanonicalizerPass());
if (job_wrapper.IsLastIRPass() && std::getenv("ONEFLOW_MLIR_ENABLE_CODEGEN_FUSERS") != nullptr) {
pm.addPass(oneflow::createOutlineJitFunctionPass());
}
...
}
Однако в этом процессе еще предстоит решить две проблемы:
- Первый вопрос, как сделать Op fusion. Приведенный выше процесс выполнения JIT учитывает только непрерывное понижение, поэтому, если есть какие-то операции, которые можно интегрировать в диалект OneFlow, что следует сделать в это время? Очень просто, давайте следовать правилам DRR MLIR или использовать синтаксис TableGen для написания серии Fuse Patterns в oneflow/ir/include/OneFlow/OneFlowPatterns.td, например,bia_add+gelu Эти две операции могут быть интегрированы в OneFlow fused_bias_add_gelu Op , то вы можете написать следующие правила.
def IsGPU: Constraint<CPred<"$0.getValue().equals("gpu")">, "is GPU device">;
def FusedBiasAddGeluPattern : Pat<
(
OneFlow_GeluOp : $gelu_op
(
OneFlow_BiasAddOp
$a,
$b,
$bias_add_op_name,
$bias_add_device_tag,
$bias_add_device_name,
$bias_add_scope_symbol_id,
$bias_add_hierarchy,
$axis
),
$gelu_op_name,
$gelu_device_tag,
$gelu_device_name,
$gelu_scope_symbol_id,
$gelu_hierarchy
),
(OneFlow_FusedBiasAddGeluOp $a, $b,
$gelu_op_name,
$gelu_device_tag,
$gelu_device_name,
$gelu_scope_symbol_id,
$gelu_hierarchy,
$axis
),
[
(IsGPU $bias_add_device_tag),
(IsGPU $gelu_device_tag)
]
>;
Здесь сопоставление выражений и перезапись выполняются на основе правил DRR MLIR.Можно видеть, что если текущим работающим устройством является графический процессор, а двумя операциями являются gelu иbias_add соответственно, они объединяются в fused_bias_add_gelu_op, что может уменьшить чтение и писать на CUDA для повышения эффективности выполнения.
- Второй вопрос: как добиться большей оптимизации некоторых операций OneFlow в инфраструктуре MLIR? Когда многоуровневый диалект спускается слой за слоем, можно увидеть, что каждая подфункция выражения MLIR OneFlow будет понижена. Впервые он будет ниже диалекта тоса.В настоящее время, если операция в этой подфункции не определяет метод преобразования в диалект тоса, то он не может быть ниже диалекта тоса. Естественно, его нельзя далее свести к Linalg Dialect, и он не может пользоваться оптимизацией, вызванной некоторыми циклическими изменениями (я чувствую, что это можно сравнить с оптимизацией планировщика TVM).
Для того, чтобы решить эту ситуацию, нам нужно определить дополнительный Pass для извлечения текущего Op или шаблона, который необходимо преобразовать в Tosa в функцию, Oneflow op в нем можно понизить до tosa, а затем oneflow mlir jit op генерируется для вызова этой функции:
def IsNotNestedInJit: Constraint<CPred<"(!$0.getDefiningOp()->getParentOfType<::mlir::FuncOp>()->hasAttr("llvm.emit_c_interface"))">, "">;
def OutlineMulCast : NativeCodeCall<"::mlir::oneflow::OutlineMulCast($_builder, $0, $1)">;
// TODO: remove attr binding if possible
def MulCastPattern : Pat<
(
OneFlow_ScalarMulByTensorOp : $mul_op
(
OneFlow_CastOp : $cast_op
$cast_x,
$cast_op_name,
$cast_device_tag,
$cast_device_name,
$cast_scope_symbol_id,
$cast_hierarchy,
$cast_dtype
),
$scalar,
$mul_op_name,
$mul_device_tag,
$mul_device_name,
$mul_scope_symbol_id,
$mul_hierarchy
),
(OutlineMulCast $mul_op, $cast_op),
[
(IsNotNestedInJit $mul_op)
]
>;
::llvm::SmallVector<::mlir::Value, 4> OutlineMulCast(::mlir::PatternRewriter& rewriter,
mlir::OpResult mul_res,
mlir::OpResult cast_res) {
if (auto mul_op = llvm::dyn_cast<ScalarMulByTensorOp>(mul_res.getDefiningOp())) {
if (auto cast_op = llvm::dyn_cast<CastOp>(cast_res.getDefiningOp())) {
// TODO: extract a function to generate op name for jit op from ops being fused
SmallString<64> op_name_storage;
auto op_name =
(cast_op.op_name() + "__FUSE__" + mul_op.op_name()).toStringRef(op_name_storage);
SmallVector<::mlir::Value, 2> operands;
operands.push_back(cast_op.in());
operands.push_back(mul_op.scalar());
SmallVector<::mlir::Value, 1> results;
results.push_back(mul_op.y());
NamedAttrList attributes =
GetJitOpAttributes(rewriter, op_name, operands.size(), results.size(), mul_op);
SmallVector<Operation*, 4> ops = {cast_op, mul_op};
auto function =
GetOrInsertFuncOp(rewriter, mul_op->getLoc(), op_name, operands, results, ops);
auto created = rewriter.create<MlirJitOp>(mul_op.getLoc(), function, attributes, operands);
assert(DumpAssembly(rewriter, created).succeeded());
cast_op->dropAllUses();
cast_op.erase();
return created->getResults();
}
}
return {};
}
void populateFuserPasses(::mlir::RewritePatternSet& patterns) {
patterns.add<MulCastPattern>(patterns.getContext());
}
Здесь нужно вручную реализовать преобразование из диалекта OneFlow в диалект Tosa шаблона MulCast и, наконец, добавить этот проход в процесс оптимизации, чтобы завершить шаблон в выражении MLIR.Он будет проходить через два уровня диалекта, Tosa и Linalg, получить некоторые возможности оптимизации.
4
Суммировать
Здесь мы берем OneFlow в качестве примера, чтобы объяснить некоторые из реальных запущенных процессов MLIR, то есть, как использовать MLIR для выполнения графа вычислений фреймворка глубокого обучения и ускорения его, В настоящее время неизбежно есть некоторые неточности в понимании , и вы можете критиковать и исправлять.
Добро пожаловать, чтобы загрузить и испытать новое поколение среды глубокого обучения OneFlow с открытым исходным кодом:GitHub.com/oneflow-Inc…