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

0x0. 背景

繼續深度學習編譯器的優化工作解讀，本篇文章要介紹的是OneFlow系統中如何基于MLIR實現Layerout Transform。在2D卷積神經網絡中，除了NCHW數據格式之外一般還存在NHWC的數據格式，對于卷積操作來說使用NHWC格式進行計算可能會獲得更好的性能。

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

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

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

開啟nn.Graph的AMP選項。

網絡選取ResNet101，對其做前向推理。

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

0x1. 實現解析

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

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

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

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

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

我們想讓其它的任意 Op 支持 NHWC 的運算，只需要定義這個接口并且重寫這個接口的成員函數即可，接下來我們看一下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 值, 表示當前的 Operator 在什么條件下是處理輸入為 NCHW 格式的數據。

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

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

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

接下來我們看一下Conv2D Op對應的 NCHWCompatibleInterface 接口實現：

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 格式的結果值。它通過創建一個新的 Conv2DOp Operation, 并將 data_format 屬性設置為 channels_last, 來實現從 NCHW 到 NHWC 的轉換。

0x1.2 插入Transpose算子

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

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

這里還涉及到幾個相關的工具函數，我們也解釋一下：

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 的轉換順序。getInputOperandTransposeOp 方法為 Operation 的輸入創建一個Transpose Operation。它使用輸入值、Transpose屬性 和 重寫器創建一個 TransposeOp , 并返回其結果。

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

0x1.3 消除多余的Transpose對

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

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的順序是否會導致冗余。它通過比較兩個 Transpose 的 perm 屬性來判斷。類似 AutoNhwcPattern ，AutoNhwcEliminateRedundantTransposePattern 類繼承自 OpRewritePattern 。它對TransposeOp 進行重寫以實現 Transpose 消除。如果順序是 NHWC->NCHW->NHWC 或NCHW->NHWC->NCHW , 則判定為冗余 Transpose 。

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

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

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));
}
};

補充：0x1.4 weight的transpose消除

這里還需要粗略的說明一下對于 weight 的 transpose 是如何處理的。在0x1.2中我們為 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 的實現。

0x2. 結論

本文介紹了一下OneFlow的編譯器中的 Layerout Transform，這個技術在后來 OneFlow 版本的 Stable Diffusion 中也發揮了重要作用，提升了推理速度。在 TVM 的 Ansor 中也有類似的優化，通過將不同的 Layerout 設定為 Op 的 strategy 進而影響 Op 的 schedule，在搜索的時候考慮到 Layerout Transform 來獲得更大的搜索空間和更好的結果。在處理Transpose 額外開銷的方法并不是唯一的，這里只是采用了一種個人認為比較簡單的方式，讀者們如果有類似需要可以自由發揮。

審核編輯：劉清