深度學(xué)習(xí)編譯器之Layerout Transform優(yōu)化

在本文的描述中，存在一些接口和Interface的混用，這兩個是一樣的都表示MLIR的Interface。

0x0. 背景

繼續(xù)深度學(xué)習(xí)編譯器的優(yōu)化工作解讀，本篇文章要介紹的是OneFlow系統(tǒng)中如何基于MLIR實現(xiàn)Layerout Transform。在2D卷積神經(jīng)網(wǎng)絡(luò)中，除了NCHW數(shù)據(jù)格式之外一般還存在NHWC的數(shù)據(jù)格式，對于卷積操作來說使用NHWC格式進(jìn)行計算可能會獲得更好的性能。

但深度學(xué)習(xí)網(wǎng)絡(luò)的訓(xùn)練一般來說是采用NCHW進(jìn)行的，我們一般只有在推理時才做NCHW到NHWC的Layerout Transform。這里存在兩個問題：首先對于一個算子比如Conv2D，它以NCHW方式訓(xùn)練時保存的權(quán)重格式是[out_channels, in_channels, *kernel_size]，但是要以NHWC格式進(jìn)行推理時我們需要對權(quán)重的格式進(jìn)行轉(zhuǎn)換；然后對于沒有權(quán)重的算子來說，我們也需要盡量的讓算子支持NHWC的運算，來減少因為卷積算子前后插入的Transpose操作帶來的額外開銷。舉個例子，假設(shè)有如下的一個小網(wǎng)絡(luò) x->conv->relu->conv->relu->out，如果我們要以NHWC格式執(zhí)行那么我們除了對2個卷積的權(quán)重進(jìn)行改動之外，我們還需要在conv前后插入transpose來修改輸入到conv算子的數(shù)據(jù)格式，也就是x->transpose(0, 2, 3, 1)->conv->transpose(0, 3, 1, 2) -> relu -> transpose(0, 2, 3, 1)->conv->transpose(0, 3, 1, 2) -> relu->out。然后細(xì)心的讀者可以發(fā)現(xiàn)，實際上這里存在很多冗余的Transpose，因為ReLU是支持以NHWC格式進(jìn)行運算的，那么這個網(wǎng)絡(luò)可以化簡為x->transpose(0, 2, 3, 1)->conv->relu->conv->relu->transpose(0, 3, 1, 2)->out。這樣可以減少一半的Transpose Op開銷。

之所以要做transpose的化簡是因為transpose算子本身也有運行以及調(diào)度的開銷，如果我們不盡量減少transpose的個數(shù)，那么因為改用NHWC帶來的計算加速可能會被 Transpose 的開銷掩蓋住。我們基于OneFlow實現(xiàn)了上述的Layerout Transform優(yōu)化，以下給出測試結(jié)果。

在V100上對這個優(yōu)化進(jìn)行了測試，測試代碼見 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/ir/test/OneFlow/auto_nhwc/test_resnet101_benchmark.py ，性能結(jié)果如下：

開啟nn.Graph的AMP選項。

網(wǎng)絡(luò)選取ResNet101，對其做前向推理。

在BatchSize=64時得到了13.6%的加速，隨著BatchSize減少加速比會減小，但始終會保持一些加速。需要注意的是，這里對權(quán)重參數(shù)部分提前進(jìn)行了transpose，所以這部分是沒有額外開銷的。實際上，我們采用了常量折疊的方式來完成，這個下篇文章再講。

0x1. 實現(xiàn)解析

在實現(xiàn)上主要需要搞定3個問題，第一個是如何確定哪些算子支持NHWC的運算，第二個是插入Transpose算子，第三個是消除多余的Transpose對。

0x1.1 基于Interface確定哪些算子支持NHWC運算

在OneFlow中如果我們要讓某個Op支持NHWC的計算，只需在Op定義時聲明一個NCHWCompatibleInterface。以卷積為例：

defOneFlow_Conv2DOp:OneFlow_ConvolutionBaseOp<"conv2d",?[NoMemoryEffect,?AttrSizedOperandSegments,?DeclareOpInterfaceMethods,DeclareOpInterfaceMethods]>{}

這里的 DeclareOpInterfaceMethods 表示這個 Operator 實現(xiàn)了 NCHWCompatibleInterface 接口,該接口定義了與 NCHW 格式兼容的 Operator 需要實現(xiàn)的方法。

我們想讓其它的任意 Op 支持 NHWC 的運算，只需要定義這個接口并且重寫這個接口的成員函數(shù)即可，接下來我們看一下NCHWCompatibleInterface 的定義。

defNCHWCompatibleInterface:OpInterface<"NCHWCompatible">{
letdescription=[{
InterfaceofNCHWcompatibility
}];

letmethods=[
InterfaceMethod<"",
????????"bool",?"IsNCHW",?(ins)
????>,
InterfaceMethod<"Create?NHWC?op?and?return?the?new?op's?results?to?be?transposed",
????????"llvm::SmallVector","NchwToNhwc",(ins"llvm::SmallVector":$transposed_inputs,"PatternRewriter&":$rewriter)
>,
InterfaceMethod<"",
????????"llvm::DenseSet","OperandsToTranspose",(ins)
>,
InterfaceMethod<"",
????????"llvm::DenseSet","ResultsToTranspose",(ins)
>,
];
letcppNamespace="::oneflow";
}

這個接口繼承自 OpInterface 接口, OpInterface 是 MLIR 框架中描述 Operator Interface 的基類。NCHWCompatibleInterface 表示一個與 NCHW 格式兼容的 Operator Interface。NCHWCompatibleInterface定義了幾個方法：

IsNCHW: 返回一個 bool 值, 表示當(dāng)前的 Operator 在什么條件下是處理輸入為 NCHW 格式的數(shù)據(jù)。

NchwToNhwc: 接受 Transpose 后的輸入和重寫器 (rewriter), 用于從 NCHW 格式轉(zhuǎn)換為 NHWC 格式。

OperandsToTranspose: 返回需要 Transpose 的輸入值集合。

ResultsToTranspose:返回需要 Transpose 的輸出值集合。

接下來我們看一下Conv2D Op對應(yīng)的 NCHWCompatibleInterface 接口實現(xiàn)：

boolConv2DOp::IsNCHW(){returnthis->getDataFormat().str()=="channels_first";}

llvm::DenseSetConv2DOp::OperandsToTranspose(){
if(this->get_addToOutput()){
return{this->getIn(),this->getWeight(),this->get_addToOutput()};
}else{
return{this->getIn(),this->getWeight()};
}
}

llvm::DenseSetConv2DOp::ResultsToTranspose(){return{this->getOut()};}

llvm::SmallVectorConv2DOp::NchwToNhwc(llvm::SmallVectorvalue,
PatternRewriter&rewriter){
autoconv_op=*this;
SmallVectoroperands;
operands.push_back(value[0]);
operands.push_back(value[1]);
if(conv_op.getBias())operands.push_back(conv_op.getBias());
if(this->get_addToOutput()){operands.push_back(value[2]);}
NamedAttrListattributes=conv_op->getAttrs();
attributes.set(conv_op.getDataFormatAttrName(),rewriter.getStringAttr("channels_last"));
autores=rewriter
.create(conv_op.getLoc(),getNHWCResultTypes(conv_op),operands,
attributes)
->getResults();
llvm::SmallVectorresults;
results.push_back(res[0]);
returnresults;
}

其中，IsNCHW 方法返回一個 bool 值,表示該 Conv2DOp Operation 是否使用 NCHW 格式。它通過檢查 Operation 的data_format 屬性來判斷。OperandsToTranspose 方法返回需要 Transpose 的輸入值集合。對于 Conv2DOp 來說,主要輸入包括input、weight、bias(可選) 和 addto_output(可選)，其中bias不需要 Transpose，并且這個addto_output是OneFlow的一個特殊的輸出用來做算子融合讀者可以忽略。

ResultsToTranspose 方法返回需要 Transpose 的輸出值集合。對于 Conv2DOp 來說,僅有一個輸出, 所以返回輸出特征圖的值。NchwToNhwc 方法接受 NCHW 格式的輸入值和重寫器,并返回 NHWC 格式的結(jié)果值。它通過創(chuàng)建一個新的 Conv2DOp Operation, 并將 data_format 屬性設(shè)置為 channels_last, 來實現(xiàn)從 NCHW 到 NHWC 的轉(zhuǎn)換。

0x1.2 插入Transpose算子

接下來就是貪心的給網(wǎng)絡(luò)里的算子插入Transpose算子，這里的思路是我們盡可能的對網(wǎng)絡(luò)里面的所有算子都前后分別插入一個Transpose，這樣的話在消除Transopose對的時候才能獲得最優(yōu)的解。給網(wǎng)絡(luò)中的算子插入Transpose的邏輯如下面的Pattern代碼所述：

structAutoNhwcPattern:publicOpInterfaceRewritePattern{
explicitAutoNhwcPattern(mlir::MLIRContext*context)
:OpInterfaceRewritePattern(context,/*benefit=*/1){}

public:
LogicalResultmatchAndRewrite(NCHWCompatibleop,PatternRewriter&rewriter)constoverride{
if(op->hasTrait()){
for(mlir::Valueoperand:op.OperandsToTranspose()){
if(operand.getType().cast().getShape().size()!=4){
returnfailure();
}
}
constautodevice_name=OpTrait::IsOpConfCompatible::getDeviceTag(op)
.cast()
.getValue()
.str();
if(device_name=="cpu"){returnfailure();}
}
llvm::SmallVectorperm=getChannelLastTransposePerm();
llvm::SmallVectorresult_perm=getChannelFirstTransposePerm();

NamedAttrListtranspose_attributes;
if(InitTransposeAttributes(op,transpose_attributes,rewriter).succeeded()){
transpose_attributes.append(llvm::StringRef("perm"),getSI32ArrayAttr(rewriter,perm));
}else{
returnfailure();
}
//whenopophasnosenseofdata_formatandpreopistranspose,wegreedilyinserttranspose
//intothisop,seekingmoreopportunitiestoeliminatetransposepattern.
constboolgreedily_transpose_flag=!op.IsNCHW()&&IsInsertTransposeOpBefore(op,rewriter);

if(op.IsNCHW()||greedily_transpose_flag){
//createtransposeopforinputoperand
SmallVectortranposed_operands;
llvm::DenseSetoperand_transpose=op.OperandsToTranspose();
intnum_transposed_operand=0;
for(Valueoperand:op->getOperands()){
if(operand_transpose.find(operand)!=operand_transpose.end()){
SmallVectorinput_res=getInputOperandTransposeOp(
op,operand,transpose_attributes,num_transposed_operand,rewriter);
tranposed_operands.push_back(input_res[0]);
num_transposed_operand+=1;
}
}
//createNHWCop
SmallVectorcreated_results=op.NchwToNhwc(tranposed_operands,rewriter);
//createtransposeopforresults
intnum_transposed_result=0;
transpose_attributes.set(llvm::StringRef("perm"),getSI32ArrayAttr(rewriter,result_perm));
llvm::DenseSettranspose_result=op.ResultsToTranspose();

for(Valueresult:op->getOpResults()){
if(transpose_result.find(result)!=transpose_result.end()){
if(autoresult_transpose_op=
getResultTransposeOp(op,created_results[num_transposed_result],
transpose_attributes,num_transposed_result,rewriter)){
result.replaceAllUsesWith(result_transpose_op);
num_transposed_result+=1;
}else{
returnfailure();
}
}
}
}
returnsuccess();
}
};

首先 AutoNhwcPattern 類繼承自 OpInterfaceRewritePattern，OpInterfaceRewritePattern 是一個用于重寫 Operation 的基類。AutoNhwcPattern 針對實現(xiàn)了 NCHWCompatible Interface 的 Operation 進(jìn)行重寫,以實現(xiàn) NCHW 到 NHWC 的格式轉(zhuǎn)換。然后，AutoNhwcPattern 重寫了 matchAndRewrite 方法。該方法會在遇到 NCHWCompatible Interface 的 Operation 時被調(diào)用,來實現(xiàn)從 NCHW 到 NHWC 的轉(zhuǎn)換。接下來，matchAndRewrite 方法首先會檢查 Operation 是否滿足轉(zhuǎn)換條件,如是否 4 維、是否在 CPU 設(shè)備上等。

如果不滿足則返回 failure。如果滿足, matchAndRewrite 方法會獲取 NCHW 到NHWC 和 NHWC 到 NCHW 的轉(zhuǎn)換順序。并初始化 Transpose Operation 的屬性。然后對于當(dāng)前 Op 是 NCHW 格式或者這個 Op 的前一個 Op 是Transpose Op，這里都進(jìn)行插入 Transpose Op的操作來獲得更多的優(yōu)化機(jī)會。

這里還涉及到幾個相關(guān)的工具函數(shù)，我們也解釋一下：

llvm::SmallVectorgetChannelLastTransposePerm(){return{0,2,3,1};}

llvm::SmallVectorgetChannelFirstTransposePerm(){return{0,3,1,2};}

llvm::SmallVectorgetInputOperandTransposeOp(NCHWCompatibleop,Valueval,
NamedAttrListtranspose_attributes,
intnum_transposed_operand,
PatternRewriter&rewriter){
std::stringtranspose_name=OpTrait::IsOpConfCompatible::getOpName(op).str()
+"_transpose_input_"+std::to_string(num_transposed_operand);
transpose_attributes.set(llvm::IsOpConfCompatible::getOpNameAttr()),
rewriter.getStringAttr(transpose_name));
SmallVectorinput_operands;
input_operands.push_back(val);
autores=rewriter
.create(op.getLoc(),getNHWCType(val.getType()),
input_operands,transpose_attributes)
->getResults();
returnres;
}

TransposeOpgetResultTransposeOp(NCHWCompatibleop,Valueval,NamedAttrListtranspose_attributes,
intnum_transposed_result,PatternRewriter&rewriter){
std::stringtranspose_name=OpTrait::IsOpConfCompatible::getOpName(op).str()
+"_transpose_output_"+std::to_string(num_transposed_result);
transpose_attributes.set(llvm::IsOpConfCompatible::getOpNameAttr()),
rewriter.getStringAttr(transpose_name));
SmallVectoroperands;
operands.push_back(val);
TransposeOptranspose_op=rewriter.create(
op.getLoc(),getNCHWType(val.getType()),operands,transpose_attributes);
returntranspose_op;
}

boolIsInsertTransposeOpBefore(NCHWCompatibleop,PatternRewriter&rewriter){
boolinsert_transpose_op_flag=false;
for(mlir::Valueoperand:op->getOperands()){
TransposeOptransposeInputOp=operand.getDefiningOp();
if(!transposeInputOp)continue;
constautoperm=transposeInputOp.getPermAttr();
if(perm.size()==4&&perm[0]==rewriter.getSI32IntegerAttr(0)
&&perm[1]==rewriter.getSI32IntegerAttr(3)&&perm[2]==rewriter.getSI32IntegerAttr(1)
&&perm[3]==rewriter.getSI32IntegerAttr(2)){
insert_transpose_op_flag=true;
break;
}
}
returninsert_transpose_op_flag;
}

其中 getChannelLastTransposePerm 和 getChannelFirstTransposePerm 方法分別返回 NHWC 到 NCHW 和 NCHW 到NHWC 的轉(zhuǎn)換順序。getInputOperandTransposeOp 方法為 Operation 的輸入創(chuàng)建一個Transpose Operation。它使用輸入值、Transpose屬性 和 重寫器創(chuàng)建一個 TransposeOp , 并返回其結(jié)果。

類似的，getResultTransposeOp 方法為 Operation 的輸出創(chuàng)建一個Transpose Operation。它使用輸出值、Transpose屬性和重寫器創(chuàng)建一個TransposeOp,并返回該Operation。IsInsertTransposeOpBefore方法檢查Operation的輸入是否已有 Transpose Operation。如果有,并且該 Transpose Operation 將 NHWC 轉(zhuǎn)為 NCHW, 則返回 true, 否則返回false。

0x1.3 消除多余的Transpose對

接下來，我們需要把插入Transpose Op的圖中所有相鄰的Transpose對盡可能的消除，代碼實現(xiàn)如下：

boolIsRedundantTransposeMatch(ArrayAttrpre,ArrayAttrafe,mlir::PatternRewriter&rewriter){
constautoprePerm=pre.getValue().vec();
constautoafePerm=afe.getValue().vec();
if(prePerm.size()==4&&afePerm.size()==4){
//handlenchw->nhwc->nchw:(0,2,3,1)->(0,3,1,2)
if(prePerm[0]==afePerm[0]&&prePerm[1]==afePerm[3]&&prePerm[2]==afePerm[1]
&&prePerm[3]==afePerm[2]&&prePerm[0]==rewriter.getSI32IntegerAttr(0)
&&prePerm[1]==rewriter.getSI32IntegerAttr(2)
&&prePerm[2]==rewriter.getSI32IntegerAttr(3)
&&prePerm[3]==rewriter.getSI32IntegerAttr(1))
returntrue;
//handlenhwc->nchw->nhwc:(0,3,1,2)->(0,2,3,1)
if(prePerm[0]==afePerm[0]&&prePerm[1]==afePerm[2]&&prePerm[2]==afePerm[3]
&&prePerm[3]==afePerm[1]&&prePerm[0]==rewriter.getSI32IntegerAttr(0)
&&prePerm[1]==rewriter.getSI32IntegerAttr(3)
&&prePerm[2]==rewriter.getSI32IntegerAttr(1)
&&prePerm[3]==rewriter.getSI32IntegerAttr(2))
returntrue;
}
returnfalse;
}

structAutoNhwcEliminateRedundantTransposePattern:publicmlir::OpRewritePattern{
explicitAutoNhwcEliminateRedundantTransposePattern(mlir::MLIRContext*context)
:OpRewritePattern(context,/*benefit=*/1){}
mlir::LogicalResultmatchAndRewrite(TransposeOpop,
mlir::PatternRewriter&rewriter)constoverride{
mlir::ValuetransposeInput=op.getOperand();
TransposeOptransposeInputOp=transposeInput.getDefiningOp();

if(!transposeInputOp
||!IsRedundantTransposeMatch(op.getPermAttr(),transposeInputOp.getPermAttr(),rewriter)){
returnfailure();
}
rewriter.replaceOp(op,{transposeInputOp.getOperand()});
returnsuccess();
}
};

IsRedundantTransposeMatch 方法檢查兩個 Transpose Operation的順序是否會導(dǎo)致冗余。它通過比較兩個 Transpose 的 perm 屬性來判斷。類似 AutoNhwcPattern ，AutoNhwcEliminateRedundantTransposePattern 類繼承自 OpRewritePattern 。它對TransposeOp 進(jìn)行重寫以實現(xiàn) Transpose 消除。如果順序是 NHWC->NCHW->NHWC 或NCHW->NHWC->NCHW , 則判定為冗余 Transpose 。

如果輸入也來自TransposeOp且兩個 Transpose 順序?qū)е氯哂?matchAndRewrite方法會用TransposeOp的輸入替換TransposeOp。實現(xiàn) Transpose 消除。matchAndRewrite 方法首先獲取 TransposeOp 的輸入,并檢查該輸入是否也來自一個 TransposeOp。如果不是, 或兩個 Transpose 的順序不導(dǎo)致冗余, 則返回 failure。最后返回 success 表示成功消除冗余 Transpose 。

最終，上面介紹的2個Pass都被封裝到 AutoNhwcPass 中作用在 MLIR 的計算圖上完成全局優(yōu)化。從下面的代碼可以看到這個優(yōu)化只有在打開 ONEFLOW_MLIR_PREFER_NHWC 環(huán)境變量時才正常生效。

voidpopulateAutoNhwcPatterns(::RewritePatternSet&patterns){
boolenable_nhwc=::ParseBooleanFromEnv("ONEFLOW_MLIR_PREFER_NHWC",false);
if(enable_nhwc){
patterns.add(patterns.getContext());
patterns.add(patterns.getContext());
}
}

classAutoNhwcPass:publicAutoNhwcPassBase{
voidrunOnOperation()override{
Operation*op=getOperation();
RewritePatternSetpatterns(op->getContext());
oneflow::populateAutoNhwcPatterns(patterns);
(void)applyPatternsAndFoldGreedily(op,std::move(patterns));
}
};

補(bǔ)充：0x1.4 weight的transpose消除

這里還需要粗略的說明一下對于 weight 的 transpose 是如何處理的。在0x1.2中我們?yōu)?weight（常量constant op） 也插入了 Transpose Op，然后我們知道 weight 是常量，所以針對 weight 的 Transpose Op 完全可以在編譯期折疊起來。這個過程是在 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/ir/oneflow-translate/lib/OneFlow/MLIROneFlowTranslation.cpp#L808-L811 這里完成的，我們后面會單獨介紹一下 Constant Folding 的實現(xiàn)。

0x2. 結(jié)論

本文介紹了一下OneFlow的編譯器中的 Layerout Transform，這個技術(shù)在后來 OneFlow 版本的 Stable Diffusion 中也發(fā)揮了重要作用，提升了推理速度。在 TVM 的 Ansor 中也有類似的優(yōu)化，通過將不同的 Layerout 設(shè)定為 Op 的 strategy 進(jìn)而影響 Op 的 schedule，在搜索的時候考慮到 Layerout Transform 來獲得更大的搜索空間和更好的結(jié)果。在處理Transpose 額外開銷的方法并不是唯一的，這里只是采用了一種個人認(rèn)為比較簡單的方式，讀者們?nèi)绻蓄愃菩枰梢宰杂砂l(fā)揮。

審核編輯：劉清