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

0x0. 背景

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

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

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

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

開啟nn.Graph的AMP選項(xiàng)。

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

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

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

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

0x1.1 基于Interface確定哪些算子支持NHWC運(yùn)算

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

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

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

我們想讓其它的任意 Op 支持 NHWC 的運(yùn)算，只需要定義這個(gè)接口并且重寫這個(gè)接口的成員函數(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";
}

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

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

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

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

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

接下來我們看一下Conv2D Op對(duì)應(yīng)的 NCHWCompatibleInterface 接口實(shí)現(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 方法返回一個(gè) bool 值,表示該 Conv2DOp Operation 是否使用 NCHW 格式。它通過檢查 Operation 的data_format 屬性來判斷。OperandsToTranspose 方法返回需要 Transpose 的輸入值集合。對(duì)于 Conv2DOp 來說,主要輸入包括input、weight、bias(可選) 和 addto_output(可選)，其中bias不需要 Transpose，并且這個(gè)addto_output是OneFlow的一個(gè)特殊的輸出用來做算子融合讀者可以忽略。

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

0x1.2 插入Transpose算子

接下來就是貪心的給網(wǎng)絡(luò)里的算子插入Transpose算子，這里的思路是我們盡可能的對(duì)網(wǎng)絡(luò)里面的所有算子都前后分別插入一個(gè)Transpose，這樣的話在消除Transopose對(duì)的時(shí)候才能獲得最優(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 是一個(gè)用于重寫 Operation 的基類。AutoNhwcPattern 針對(duì)實(shí)現(xiàn)了 NCHWCompatible Interface 的 Operation 進(jìn)行重寫,以實(shí)現(xiàn) NCHW 到 NHWC 的格式轉(zhuǎn)換。然后，AutoNhwcPattern 重寫了 matchAndRewrite 方法。該方法會(huì)在遇到 NCHWCompatible Interface 的 Operation 時(shí)被調(diào)用,來實(shí)現(xiàn)從 NCHW 到 NHWC 的轉(zhuǎn)換。接下來，matchAndRewrite 方法首先會(huì)檢查 Operation 是否滿足轉(zhuǎn)換條件,如是否 4 維、是否在 CPU 設(shè)備上等。

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

這里還涉及到幾個(gè)相關(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)建一個(gè)Transpose Operation。它使用輸入值、Transpose屬性 和 重寫器創(chuàng)建一個(gè) TransposeOp , 并返回其結(jié)果。

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

0x1.3 消除多余的Transpose對(duì)

接下來，我們需要把插入Transpose Op的圖中所有相鄰的Transpose對(duì)盡可能的消除，代碼實(shí)現(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 方法檢查兩個(gè) Transpose Operation的順序是否會(huì)導(dǎo)致冗余。它通過比較兩個(gè) Transpose 的 perm 屬性來判斷。類似 AutoNhwcPattern ，AutoNhwcEliminateRedundantTransposePattern 類繼承自 OpRewritePattern 。它對(duì)TransposeOp 進(jìn)行重寫以實(shí)現(xiàn) Transpose 消除。如果順序是 NHWC->NCHW->NHWC 或NCHW->NHWC->NCHW , 則判定為冗余 Transpose 。

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

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

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消除

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

0x2. 結(jié)論

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

審核編輯：劉清