以O(shè)neFlow為例探索MLIR的實(shí)際開發(fā)流程
前言
最近在同事shenghang的幫助下做了一點(diǎn)OneFlow IR相關(guān)的開發(fā)，對MLIR執(zhí)行部分有一些新的感受，所以嘗試分享一下。我之前花了不少時(shí)間去理解OneFlow IR的整個(gè)架構(gòu)（可以看我的Toy Tutorials系列），但對OneFloiw IR的JIT的執(zhí)行這部分一直存疑。最近將OneFlow基于Job（OneFlow的作業(yè)函數(shù)，不考慮設(shè)備的話可以理解為一個(gè)計(jì)算圖）接入MLIR工程實(shí)現(xiàn)部分重新進(jìn)行了梳理，并在shenghang的指導(dǎo)下理解了整個(gè)流程。所以這篇文檔我將介紹一下OneFlow和MLIR是如何結(jié)合的，如何在OneFlow IR中新增一個(gè)圖級別的Pass，OneFlow的Operation是如何自動(dòng)變成MLIR 的Operation的以及為什么OneFlow IR能利用MLIR為計(jì)算帶來加速等。我對MLIR的了解不算多，2個(gè)月前開始接觸，有任何錯(cuò)誤請大家批評斧正。本文和 https://github.com/Oneflow-Inc/oneflow & https://github.com/BBuf/tvm_mlir_learn 有關(guān)，感興趣可以star關(guān)注一下。
本文提到的Op和Operation是一回事，沒有嚴(yán)格區(qū)分。
OneFlow是如何和MLIR結(jié)合的？
在OneFlow中引入MLIR作為OneFlow的IR有諸多優(yōu)點(diǎn)，不僅可以取代OneFlow中需要通過C++手寫的Operation定義減小開發(fā)難度，還可以降低Operation定義中一些容器相關(guān)的開銷。另外我們還可以通過MLIR維護(hù)的基礎(chǔ)設(shè)施（即多重Dialect）來完成對計(jì)算圖計(jì)算的加速。這里的計(jì)算圖既可以是Eager的計(jì)算圖，也可以是Lazy的計(jì)算圖。由于基于Eager計(jì)算圖使用MLIR進(jìn)行加速的工作（即oneflow.jit.xxx）還沒有正式開放，我這里仍然以Lazy計(jì)算圖（Job）為例來講解OneFlow和MLIR的結(jié)合過程。
首先我們需要編譯好開啟MLIR的OneFlow，編譯命令如下：
[email protected]:Oneflow-Inc/oneflow.git
cd?oneflow?&&?mkdir?build?&&?cd?build
cmake-C?../cmake/caches/cn/fast/mlir-cuda-75.cmake?-DBUILD_TESTING=ON?..?&&?ninja?
然后可以寫一個(gè)例子進(jìn)行測試：
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()))
運(yùn)行這個(gè)例子之后會在當(dāng)前運(yùn)行目錄下生成一個(gè)log文件，里面有一個(gè)ir_pass 文件夾記錄了經(jīng)過OneFlow MLIR優(yōu)化前后的計(jì)算圖(.prototxt) 以及 MLIR的表達(dá)式(*.mlir)，還有一個(gè)*.mlir.dot文件可以用graphviz打開來可視化MLIR表達(dá)式的計(jì)算圖。需要注意的是如果OneFlow正在執(zhí)行訓(xùn)練任務(wù)，這個(gè)log文件夾里不僅包含前向的計(jì)算圖和MLIR表達(dá)式，也會生成后向的計(jì)算圖和MLIR表達(dá)式。所以MLIR在整個(gè)神經(jīng)網(wǎng)絡(luò)的運(yùn)行流程中均可以作用，這是區(qū)別于前向推理框架的重要一點(diǎn)，即訓(xùn)練也可以加速。
在oneflow/api/python/ir.cpp 中有下面兩行代碼：
REGISTER_JOB_PASS("IRRoundTripBeforeAD",?IRRoundTrip);
REGISTER_JOB_PASS("IRRoundTrip",?IRRoundTrip);
RoundTrip即往返的意思，BeforeAD可以理解為反向之前，kAfterAD 可以理解為反向之后，這里通過將OneFlow Job和MLIR的互轉(zhuǎn)過程注冊為OneFlow Job的一個(gè)Pass來建立OneFlow計(jì)算圖和MLIR的聯(lián)系。在執(zhí)行OneFlow腳本時(shí)，如果想使能MLIR作用于OneFlow計(jì)算圖，開啟ONEFLOW_MLIR_ENABLE_ROUND_TRIP=1環(huán)境變量即可。
接下來，要將OneFlow的計(jì)算圖和MLIR建立聯(lián)系等價(jià)于將OneFlow計(jì)算圖中的Operation和MLIR中的Operation進(jìn)行一對一的轉(zhuǎn)換。而MLIR的Operation定義在各級Dialect下，按照MLIR的通用接入原則，我們實(shí)現(xiàn)了一個(gè)OneFlow Dialect并在OneFlow Dialect上實(shí)現(xiàn)了OneFlow Operation到OneFlow Dialect下的Operation的一一映射。如何定義OneFlow Dialect和Operation這里就不講了，可以參考MLIR官方文檔的Dialects和ODS一節(jié)（https://mlir.llvm.org/docs/OpDefinitions/）或者我之前的文章，它們都是基于TableGen規(guī)則來完成的。關(guān)于MLIR Operation的定義我之前結(jié)合OneFlow Dialect的Op定義總結(jié)了一個(gè)文檔（https://github.com/BBuf/tvm_mlir_learn 中） 。除了Dialect和Operation的定義還有一些其它需要定義的東西，比如OneFlow數(shù)據(jù)類型到MLIR數(shù)據(jù)類型映射的定義在oneflow/ir/include/OneFlow/OneFlowEnums.td ，OneFlow Dialect Operation的一些通用前端接口定義在oneflow/ir/include/OneFlow/OneFlowEnums.td。這里我們以Reshape Operation為例子來簡單說明一下這個(gè)Operation有哪些組成部分：
def?OneFlow_ReshapeOp?:?OneFlow_BaseOp<"reshape",?[NoSideEffect,?DeclareOpInterfaceMethods]>?{
??let?input?=?(ins
????AnyType:$in
??);
??let?output?=?(outs
????AnyType:$out
??);
??let?attrs?=?(ins
????AnyI64ElementsAttr:$shape
??);
}
OneFlow_ReshapeOp 這個(gè)名字下劃線之前的是Dialect的名字，后面是這個(gè)Dialect下的Operation的名字。然后這個(gè)Operation繼承了OneFlow_BaseOp基類，并聲明了約束和前端接口，接下來定義了Operation的輸入，輸出和屬性就結(jié)束了。可以發(fā)現(xiàn)OneFlow Dialect Operation的定義和OneFlow User Op是完全一致的，這保證了OneFlow和MLIR互轉(zhuǎn)的合法性。OneFlow Reshape Operation的定義如下：
REGISTER_USER_OP("reshape")
????.Input("in")
????.Output("out")
????.Attr("shape")
????...
OneFlow Job和MLIR的互轉(zhuǎn)實(shí)現(xiàn)在oneflow/ir/oneflow-translate，主要做的事情就是遍歷Job的OpGraph，對節(jié)點(diǎn)和邊分別進(jìn)行處理最后轉(zhuǎn)換成一個(gè)MLIR表達(dá)式，同時(shí)在計(jì)算完成后可以基于MLIR表達(dá)式重寫Job。這里的整體邏輯偏復(fù)雜，因?yàn)橐幚鞳neFlow Job OpGraph里面各種類型Operation和邊的轉(zhuǎn)化，這里不繼續(xù)深入講解，因?yàn)樗膊皇俏疫@篇文章要討論的點(diǎn)，感興趣的可以直接閱讀代碼。
OneFlow IR如何執(zhí)行？
在上面Operation定義時(shí)是舉了一個(gè)Reshape的例子，瀏覽oneflow/ir/include/OneFlow/OneFlowOps.td容易發(fā)現(xiàn)這里還定義了一個(gè)OneFlow_MlirJitOp，這個(gè)自定義的Op就是用來執(zhí)行MLIR表達(dá)式的，它里面實(shí)現(xiàn)了CPU和GPU的Kernel（源碼在oneflow/ir/oneflow-extension/extension.cpp）用來加載MLIR提供的JIT執(zhí)行引擎運(yùn)行最終得到的LLVM IR。那么LLVM IR又是怎么來的呢？這是通過OneFlow MLIR表達(dá)式逐級下降之后得來的，具體下降過程如下：
void?AddLowerToLinalgMemRefPasses(PassManager&?pm)?{
??pm.addPass(createLowerOneFlowToTosaPass());????????????//?lower-oneflow-to-tosa
??pm.addPass(createCSEPass());???????????????????????????//?cse
??pm.addNestedPass(tosa::createTosaToLinalg());??//?tosa-to-linalg-on-tensors
??auto?p?=?createLinalgElementwiseOpFusionPass();
??assert(p->initializeOptions("allow-folding-unit-dim-reshapes=true").succeeded());
??pm.addNestedPass(std::move(p));?????????????????????//?linalg-fuse-elementwise-ops
??pm.addNestedPass(createLinalgBufferizePass());??????//?linalg-bufferize
??pm.addNestedPass(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(createFinalizingBufferizePass());??//?finalizing-bufferize
}
LogicalResult?LowerModuleToLLVM(mlir::MLIRContext*?context,?ModuleOp?module)?{
??mlir::PassManager?pm(context);
??AddLowerToLinalgMemRefPasses(pm);
??pm.addNestedPass(createConvertLinalgToLoopsPass());??//?convert-linalg-to-loops
??pm.addNestedPass(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 Dialect首先下降到Tosa Dialect，然后下降到Linalg Dialect，再然后是Loop Dialect，一直到最后的LLVM IR。在逐級下降的過程中，我們可以享受如Linalg Dialect帶來的嵌套循環(huán)變換帶來的優(yōu)化機(jī)會以提升最終IR的性能。這里的Lowering過程是在OneFlow調(diào)用MlirJitOp 的Kernel時(shí)觸發(fā)的（oneflow/ir/oneflow-extension/extension.cpp ），調(diào)用也是作為一個(gè)MLIR的Pass被加入到了優(yōu)化流程中。JIT調(diào)用流程Pass的實(shí)現(xiàn)可以精簡為：
class?OutlineJitFunctionPass?:?public?OutlineJitFunctionPassBase?{
??void?runOnOperation()?override?{
????Operation*?op?=?getOperation();
????RewritePatternSet?patterns(op->getContext());
????oneflow::populateFuserPasses(patterns);
????(void)applyPatternsAndFoldGreedily(op,?std::move(patterns));
??}
};
std::unique_ptr?createOutlineJitFunctionPass()? {
??return?std::make_unique();
}
LogicalResult?ApplyRoundTripPatterns(RoundTripOneFlowJobWrapperInterface&?job_wrapper,
?????????????????????????????????????MLIRContext*?context,?OwningModuleRef&?module)?{
??mlir::PassManager?pm(context);
??pm.addNestedPass(::mlir::createCanonicalizerPass());
??if?(job_wrapper.IsLastIRPass()?&&?std::getenv("ONEFLOW_MLIR_ENABLE_CODEGEN_FUSERS")?!=?nullptr)?{
????pm.addPass(oneflow::createOutlineJitFunctionPass());
??}
??...
}
但這套流程還存在兩個(gè)問題需要解決：
第一個(gè)問題是如何做Op融合。上面的JIT執(zhí)行流程只考慮了不斷Lowering，那么假如在OneFlow Dialect中有一些Operation是可以融合的，這個(gè)時(shí)候應(yīng)該怎么做呢？很簡單，我們沿用一下MLIR的DRR規(guī)則，還是用TableGen語法在 oneflow/ir/include/OneFlow/OneFlowPatterns.td中寫一系列的Fuse Pattern即可，比如bias_add+gelu這兩個(gè)Op可以融合成OneFlow中的fused_bias_add_geluOp，那么就可以寫如下的規(guī)則。
def?IsGPU:?Constraint"$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)
??]
>;
這里基于MLIR的DRR規(guī)則來做表達(dá)式匹配和重寫，可以看到假如當(dāng)前運(yùn)行設(shè)備是GPU并且前后兩個(gè)Op分別是gelu和bias_add 就將其進(jìn)行融合為一個(gè)fused_bias_add_gelu_op，在CUDA上可以減少讀寫來提升執(zhí)行效率。
第二個(gè)問題是如何讓OneFlow的一些Operation享受MLIR基礎(chǔ)設(shè)施中的更多優(yōu)化？在多級Dialect 逐層下降時(shí)可以看到OneFlow的MLIR表達(dá)式的每個(gè)子函數(shù)都會被Lower。第一次會將其Lower到Tosa Dialect，這個(gè)時(shí)候如果這個(gè)子函數(shù)中的某個(gè)Operation沒有定義轉(zhuǎn)換到Tosa Dialect的方法，那么就不能Lower到Tosa Dialect。自然也就不能進(jìn)一步下降為Linalg Dialect，享受不到一些循環(huán)變化帶來的優(yōu)化（我感覺可以類比TVM的scheduler優(yōu)化）。為了解決這種情況我們需要額外再定義一個(gè)Pass來將當(dāng)前需要轉(zhuǎn)換為Tosa的Op或者模式提取成一個(gè)函數(shù)，里面的oneflow op都能夠lower到tosa，然后生成一個(gè) oneflow mlir jit op 來 call 這個(gè)函數(shù)：
def?IsNotNestedInJit:?Constraint"(!$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(mul_res.getDefiningOp()))?{
????if?(auto?cast_op?=?llvm::dyn_cast(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);
??????SmallVector4>?ops?=?{cast_op,?mul_op};
??????auto?function?=
??????????GetOrInsertFuncOp(rewriter,?mul_op->getLoc(),?op_name,?operands,?results,?ops);
??????auto?created?=?rewriter.create(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(patterns.getContext());
}
這里就是將MulCast這個(gè)Pattern手動(dòng)實(shí)現(xiàn)了從OneFlow Dialect到Tosa Dialect的轉(zhuǎn)換，最后將這個(gè)Pass加到優(yōu)化流程中即可完成MLIR表達(dá)式中的這個(gè)Pattern會經(jīng)過Tosa和Linalg這兩個(gè)層次的Dialect，獲得一些優(yōu)化機(jī)會。
總結(jié)
這里以O(shè)neFlow為例講解了一些MLIR的真實(shí)運(yùn)行流程，即是如何通過MLIR來執(zhí)行深度學(xué)習(xí)框架的計(jì)算圖并且為其加速的，目前理解難免有不到位的地方，歡迎大家批評指正。