RM新时代|官方理财平台

<pre id="7wtpw"><big id="7wtpw"><strike id="7wtpw"></strike></big></pre>

算豐技術(shù)揭秘｜基于TPU-MLIR：詳解EinSum的完整處理過(guò)程！

發(fā)布時(shí)間: 2024年2月6日

EinSum介紹

EinSum（愛(ài)因斯坦求和）是一個(gè)功能強大的算子，能夠簡(jiǎn)潔高效地表示出多維算子的乘累加過(guò)程，對使用者非常友好。

本質(zhì)上， EinSum是一個(gè)算子族，可以表示多種基礎操作，如矩陣乘法、Reduce。EinSum支持任意多的輸入，只要計算中只包含點(diǎn)乘（element-wise）、廣播（broadcast）、歸約求和（reduction sum）都可以使用EinSum來(lái)表示。以下給出一種將EinSum計算等價(jià)表達的流程：

將輸入的維度符號放入一個(gè)列表，移除重復元素后按升序排列；
對各輸入維度執行轉置操作，確保維度標識符按照升序對齊，實(shí)現維度對齊；
在缺失的維度上填充1（擴展維度），以便與第一步中定義的維度保持一致；
對所有輸入執行廣播點(diǎn)乘；
對那些不在輸出標識符中的維度執行累加操作；
利用轉置操作調整維度順序，使其與輸出標識符的順序一致。

下圖是以out = EinSum("ijk, lki-> li", in0, in1)為例，根據上述步驟進(jìn)行等價(jià)轉換。

TPU-MLIR轉換

雖然使用上述流程可以完成對EinSum的計算轉換，但如果嚴格按照該流程執行，會(huì )帶來(lái)大量的Transpose和Reshape操作，這不僅會(huì )給TPU-MLIR的LayerGroup功能帶來(lái)挑戰，同時(shí)也難以顯式地識別出如矩陣乘法這類(lèi)操作，從而無(wú)法充分利用硬件加速單元。因此，TPU-MLIR并未直接采用上述流程進(jìn)行轉換。

接下來(lái)，我們將詳細介紹EinSum的完整處理過(guò)程。

前端接口

以下示例代碼摘自OnnxConverter.py文件，并附帶了注釋。代碼整體結構簡(jiǎn)潔明了，我們可以看到，轉換函數目前僅支持兩個(gè)輸入的常見(jiàn)情況。特別需要注意的是公式的歸一化過(guò)程。由于EinSum的表達式可以使用任意非重復字符來(lái)表示下標，這雖然提高了可讀性，但也導致同一操作有多種不同的表示方式。歸一化操作就是將表達式字符重新映射，以字符'a'作為起始。例如，比如ij,jk->ik和dk,kv->dv都會(huì )映射為ab,bc->ac。

????#?https://pytorch.org/docs/1.13/generated/torch.einsum.html?highlight=einsum#torch.einsum
????def?convert_einsum_op(self,?onnx_node):
????????assert?(onnx_node.op_type?==?"Einsum")
????????equation?=?onnx_node.attrs.get("equation").decode()

????????#?公式歸一化
????????def?normalize_equation(equation_c):
????????????equation?=?equation_c
????????????new_equation?=?''
????????????start?=?'a'
????????????translate_map?=?{}
????????????for?s?in?equation:
????????????????if?s?==?'?':
????????????????????continue
????????????????elif?not?((s?>=?'a'?and?s?<=?'z')?or?(s?>=?'A'?and?s?<=?'Z')):
????????????????????translate_map[s]?=?s
????????????????elif?s?not?in?translate_map:
????????????????????translate_map[s]?=?start
????????????????????start?=?chr(ord(start)?+?1)
????????????????new_equation?+=?translate_map[s]
????????????return?new_equation
????????equation?=?normalize_equation(equation)
????????lhs?=?self.getOperand(onnx_node.inputs[0])?#
????????#?大多情況下rhs是Weight, self.getOp會(huì )先到Weight Map中查找；如果找不到，
????????#?其會(huì )從Mutable Tensor中查找，然后返回對應的Value。
????????rhs?=?self.getOp(onnx_node.inputs[1])
????????new_op?=?top.EinsumOp(self.unranked_type,
??????????????????????????????[lhs,?rhs],
??????????????????????????????mode=StringAttr.get(equation),
??????????????????????????????#?設置?loc?信息，方便找到原圖對應算子
??????????????????????????????loc=self.get_loc("{}_{}".format(onnx_node.name,?onnx_node.op_type)),
??????????????????????????????#?將該算子插入到當前的block中
??????????????????????????????ip=self.mlir.insert_point).output
????????#?將輸出放到Mutable?Tensor列表中，供后面算子使用
????????self.addOperand(onnx_node.name,?new_op)

內部轉換

TPU-MLIR目前支持了幾種常見(jiàn)的表達式，并根據不同的算子進(jìn)行了優(yōu)化轉換。所有的變換最終都利用了硬件的矩陣乘法加速單元，從而實(shí)現了對算子的有效加速。以下是部分代碼片段，該代碼來(lái)自tpu-mlir/lib/Dialect/Top/Canonicalize/Einsum.cpp，并在原有基礎上添加了注釋。

struct?ConvertEinsum?:?public?OpRewritePattern<EinsumOp>?{
??using?OpRewritePattern::OpRewritePattern;

??LogicalResult?matchAndRewrite(EinsumOp?op,
????????????????????????????????PatternRewriter?&rewriter)?const?override?{
????//?目前只支持輸入個(gè)數為2或者輸入0為Weight的情況
????if?(op.getInputs().size()?!=?2?||?module::isWeight(op.getInputs()[0]))?{
??????llvm_unreachable("Not?support?now.");
??????//?return?failure();
????}
????auto?none?=?module::getNoneOp(op);
????auto?mode?=?op.getMode().str();
????auto?lhs?=?op.getInputs()[0];
????auto?rhs?=?op.getInputs()[1];
????auto?lshape?=?module::getShape(lhs);
????auto?rshape?=?module::getShape(rhs);
????std::string?lname?=?module::getName(lhs).str();
????std::string?rname?=?module::getName(rhs).str();
????std::string?name?=?module::getName(op.getOutput()).str();

????std::vector<Value>?operands;
????std::vector<NamedAttribute>?attrs;
????if?(mode?==?"a,b->ab")?{
??????//?外積操作:?可看作[a,1]x[1,b]的矩陣乘法操作
??????//?lhs->ReshapeOp():?shape=[a]?to?shape[a,1]
??????rewriter.setInsertionPointAfter(lhs.getDefiningOp());
??????//
??????auto?newType?=?RankedTensorType::get({lshape[0],?1},?module::getElementType(lhs));
??????auto?loc?=?NameLoc::get(rewriter.getStringAttr(lname?+?"_to2dim"));
??????auto?lrsOp?=?rewriter.create<ReshapeOp>(loc,?newType,?ValueRange{lhs});
??????operands.push_back(lrsOp);

??????//?rhs->ReshapeOp():?shape=[b]?to?shape[1,b]
??????rewriter.setInsertionPointAfter(rhs.getDefiningOp());
??????newType?=?RankedTensorType::get({1,?rshape[0]},?module::getElementType(rhs));
??????loc?=?NameLoc::get(rewriter.getStringAttr(rname?+?"_to2dim"));
??????auto?rrsop?=?rewriter.create<ReshapeOp>(loc,?newType,?ValueRange{rhs});
??????operands.push_back(rrsop);
??????operands.push_back(none);
??????//?用MatMulOp實(shí)現[a,1]x[1,b]=[a,b],?并替換原來(lái)的EinSum操作
??????rewriter.setInsertionPoint(op);
??????auto?matmulOp?=?rewriter.create<MatMulOp>(op.getLoc(),?op.getType(),?operands,?attrs);
??????op.replaceAllUsesWith(matmulOp.getOperation());
??????rewriter.eraseOp(op);
????}?else?if?(mode?==?"abcd,cde->abe")?{
??????//?可以轉換成矩陣乘法[a*b,?c*d]x[c*d,?e]->[a*b,?e]->[a,b,e]
??????//?lhs_reshape_rst?=?[lhs_shape[0]?*?lhs_shape[1],?lhs_shape[2]?*?lhs_shape[3]]
??????rewriter.setInsertionPointAfter(lhs.getDefiningOp());
??????auto?newType?=?RankedTensorType::get({lshape[0]?*?lshape[1],?lshape[2]?*?lshape[3]},?module::getElementType(lhs));
??????auto?loc?=?NameLoc::get(rewriter.getStringAttr(lname?+?"_to2dim"));
??????auto?lrsOp?=?rewriter.create<ReshapeOp>(loc,?newType,?ValueRange{lhs});
??????operands.push_back(lrsOp);
??????newType?=?RankedTensorType::get({rshape[0]?*?rshape[1],?rshape[2]},?module::getElementType(rhs));
??????if?(module::isWeight(rhs))?{
????????rhs.setType(newType);
????????operands.push_back(rhs);
??????}?else?{
????????rewriter.setInsertionPointAfter(rhs.getDefiningOp());
????????loc?=?NameLoc::get(rewriter.getStringAttr(rname?+?"_to2dim"));
????????auto?rrsop?=?rewriter.create<ReshapeOp>(loc,?newType,?ValueRange{rhs});
????????operands.push_back(rrsop);
??????}
??????operands.push_back(none);
??????rewriter.setInsertionPoint(op);
??????newType?=?RankedTensorType::get({lshape[0]?*?lshape[1],?rshape[2]},?module::getElementType(op));
??????loc?=?NameLoc::get(rewriter.getStringAttr(name?+?"_matmul"));
??????auto?matmulOp?=?rewriter.create<MatMulOp>(loc,?newType,?operands,?attrs);
??????auto?orsOp?=?rewriter.create<ReshapeOp>(op.getLoc(),?op.getType(),?ValueRange{matmulOp});
??????op.replaceAllUsesWith(orsOp.getOperation());
??????rewriter.eraseOp(op);
????}?else?if?(mode?==?"abcd,bed->abce")?{
??????rewriter.setInsertionPointAfter(rhs.getDefiningOp());
??????//?轉換過(guò)程
??????//?batch?matmul?does?not?support?broadcast
??????//?temporary?solution
??????//?[h,?k,?c]?->?[1,?h,?k,?c]?->?[b,?h,?k,?c]
??????operands.push_back(lhs);

??????RankedTensorType?newType;
??????//?右操作數處理
??????if?(auto?wOp?=?dyn_cast<top::WeightOp>(rhs.getDefiningOp()))?{
????????//?對于Weight來(lái)說(shuō)，可以將數據復制，解決不支持廣播問(wèn)題，?[b,?e,?d]->[a,?b,?e,?d]
????????auto?storage_type?=?module::getStorageType(rhs);
????????assert(storage_type.isF32()?&&?"Todo,?supoort?more?weight?type");
????????auto?data?=?wOp.read_as_byte();
????????uint8_t?*dptr;
????????newType?=?RankedTensorType::get({lshape[0],?rshape[0],?rshape[1],?rshape[2]},?module::getElementType(rhs));
????????std::vector<float_t>?new_filter(newType.getNumElements(),?0);
????????dptr?=?(uint8_t?*)new_filter.data();
????????//?實(shí)際的數據復制過(guò)程
????????for?(int32_t?i?=?0;?i?<?lshape[0];?i++)?{
??????????auto?offset?=?i?*?data->size();
??????????memcpy(dptr?+?offset,?data->data(),?data->size());
????????}
????????auto?new_op?=?top::WeightOp::create(op,?"folder",?new_filter,?newType);
????????wOp.replaceAllUsesWith(new_op.getDefiningOp());
????????operands.push_back(new_op);
????????rewriter.eraseOp(wOp);
??????}?else?{
????????//?對于普通tensor,?先reshape成[1,?b,?e,?d]?再用tile算子翻倍數據為?[a,?b,?e,?d]

????????//?Reshape操作
????????auto?loc?=?NameLoc::get(rewriter.getStringAttr(rname?+?"_reshape"));
????????newType?=?RankedTensorType::get({1,?rshape[0],?rshape[1],?rshape[2]},?module::getElementType(rhs));
????????auto?rrsop?=?rewriter.create<ReshapeOp>(loc,?newType,?ValueRange{rhs});

????????//?Tile操作，各維tile倍數[a,1,1,1]
????????newType?=?RankedTensorType::get({lshape[0],?rshape[0],?rshape[1],?rshape[2]},?module::getElementType(rhs));
????????loc?=?NameLoc::get(rewriter.getStringAttr(rname?+?"_tile"));
????????attrs.push_back(rewriter.getNamedAttr("tile",?rewriter.getI64ArrayAttr({lshape[0],?1,?1,?1})));
????????auto?tileOp?=?rewriter.create<TileOp>(loc,?newType,?ValueRange{rrsop},?attrs);
????????attrs.clear();
????????operands.push_back(tileOp);
??????}
??????operands.push_back(none);
??????//?這里使用了右操作數轉置的批量矩陣乘法算子,?硬件可直接支持
??????//?[a*b,?c,?d]?*?[a*b,?e,?d]^T?->?[a*b,?c,?e]
??????attrs.push_back(rewriter.getNamedAttr("right_transpose",?rewriter.getBoolAttr(true)));
??????rewriter.setInsertionPoint(op);
??????auto?matmulOp?=?rewriter.create<MatMulOp>(op.getLoc(),?op.getType(),?operands,?attrs);
??????op.replaceAllUsesWith(matmulOp.getOperation());
??????rewriter.eraseOp(op);
????}?else?if?(mode?==?"abcd,ced->abce")?{
??????//?dumb?implementation
??????//?轉置lhs?[a,?b,?c,?d]?->?[a,?c,?b,?d]
??????//?trans_shape?=?[lhs_shape[0],?lhs_shape[2],?lhs_shape[1],?lhs_shape[3]]
??????rewriter.setInsertionPointAfter(lhs.getDefiningOp());
??????auto?loc?=?NameLoc::get(rewriter.getStringAttr(lname?+?"_trans"));
??????auto?newType?=?RankedTensorType::get({lshape[0],?lshape[2],?lshape[1],?lshape[3]},?module::getElementType(lhs));
??????attrs.push_back(rewriter.getNamedAttr("order",?rewriter.getI64ArrayAttr({0,?2,?1,?3})));
??????auto?tranOp?=?rewriter.create<PermuteOp>(loc,?newType,?ValueRange{lhs},?attrs);
??????attrs.clear();
??????operands.push_back(tranOp);

??????//?復制或Tile?lhs:?[c,e,d]?->?[a,c,e,d]
??????rewriter.setInsertionPointAfter(rhs.getDefiningOp());
??????if?(auto?wOp?=?dyn_cast<top::WeightOp>(rhs.getDefiningOp()))?{
????????//?Weight翻倍數據
????????auto?storage_type?=?module::getStorageType(rhs);
????????assert(storage_type.isF32()?&&?"Todo,?supoort?more?weight?type");
????????auto?data?=?wOp.read_as_byte();
????????uint8_t?*dptr;
????????newType?=?RankedTensorType::get({lshape[0],?rshape[0],?rshape[1],?rshape[2]},?module::getElementType(rhs));
????????std::vector<float_t>?new_filter(newType.getNumElements(),?0);
????????dptr?=?(uint8_t?*)new_filter.data();
????????for?(int32_t?i?=?0;?i?<?lshape[0];?i++)?{
??????????auto?offset?=?i?*?data->size();
??????????memcpy(dptr?+?offset,?data->data(),?data->size());
????????}
????????auto?new_op?=?top::WeightOp::create(op,?"folder",?new_filter,?newType);
????????wOp.replaceAllUsesWith(new_op.getDefiningOp());
????????operands.push_back(new_op);
????????rewriter.eraseOp(wOp);
??????}?else?{
????????//?rehshape?+?tile:?[c,e,d]?-reshape->[1,c,e,d]-tile->[a,c,e,d]
????????loc?=?NameLoc::get(rewriter.getStringAttr(rname?+?"_reshape"));
????????newType?=?RankedTensorType::get({1,?rshape[0],?rshape[1],?rshape[2]},?module::getElementType(rhs));
????????auto?rrsop?=?rewriter.create<ReshapeOp>(loc,?newType,?ValueRange{rhs});
????????loc?=?NameLoc::get(rewriter.getStringAttr(rname?+?"_tile"));
????????attrs.push_back(rewriter.getNamedAttr("tile",?rewriter.getI64ArrayAttr({lshape[0],?1,?1,?1})));
????????newType?=?RankedTensorType::get({lshape[0],?rshape[0],?rshape[1],?rshape[2]},?module::getElementType(rhs));
????????auto?tileOp?=?rewriter.create<TileOp>(loc,?newType,?ValueRange{rrsop},?attrs);
????????attrs.clear();
????????operands.push_back(tileOp);
??????}
??????operands.push_back(none);
??????//?右操作數帶轉置批量矩陣乘法：[a*c, b, d]?*?[a*c, e, d]^T ->?[a*c, b, e]->[a, c, b, e]
??????newType?=?RankedTensorType::get({lshape[0],?lshape[2],?lshape[1],?rshape[1]},?module::getElementType(op));
??????attrs.push_back(rewriter.getNamedAttr("right_transpose",?rewriter.getBoolAttr(true)));
??????rewriter.setInsertionPoint(op);
??????loc?=?NameLoc::get(rewriter.getStringAttr(name?+?"_matmul"));
??????auto?matmulOp?=?rewriter.create<MatMulOp>(loc,?newType,?operands,?attrs);
??????attrs.clear();
??????//?[b,?w,?h,?k]?->?[b,?h,?w,?k]
??????attrs.push_back(rewriter.getNamedAttr("order",?rewriter.getI64ArrayAttr({0,?2,?1,?3})));
??????auto?tranBackOp?=?rewriter.create<PermuteOp>(op.getLoc(),?op.getType(),?ValueRange{matmulOp},?attrs);
??????op.replaceAllUsesWith(tranBackOp.getOperation());
??????rewriter.eraseOp(op);
????}?else?if?(mode?==?"abcd,abed->abce"?||?mode?==?"abcd,abde->abce")?{
??????//?lhs(abcd)?*?rhs(abed)^T?->?abce
??????//?lhs(abcd)?*?rhs(abde)?->?abce
??????auto?newType?=?RankedTensorType::get({lshape[0],?lshape[1],?lshape[2],?rshape[2]},?module::getElementType(op));
??????if?(mode?==?"abcd,abde->abce"){
????????newType?=?RankedTensorType::get({lshape[0],?lshape[1],?lshape[2],?rshape[3]},?module::getElementType(op));
??????}
??????rewriter.setInsertionPoint(op);
??????rewriter.setInsertionPointAfter(rhs.getDefiningOp());
??????operands.push_back(lhs);
??????operands.push_back(rhs);
??????operands.push_back(none);
??????if?(mode?==?"abcd,abed->abce"){
????????//rhs(abed)^T
????????attrs.push_back(rewriter.getNamedAttr("right_transpose",?rewriter.getBoolAttr(true)));
??????}

??????auto?loc?=?NameLoc::get(rewriter.getStringAttr(name));
??????auto?matmulOp?=?rewriter.create<MatMulOp>(loc,?newType,?operands,?attrs);
??????op.replaceAllUsesWith(matmulOp.getOperation());
??????attrs.clear();
??????rewriter.eraseOp(op);

????}??else?if?(mode?==?"abcd,cde->abce"){

????//?lhs?:
????//?????abcd?->?acbd(pemute)
????//?rhs?:
????//?????cde??->?1cde(reshape)
????//?????acde?->?acde(tile)
????//?matmul:
????//???lhs(acbd)?*?rhs(acde)?=?result(acbe)
????//?result:
????//?????acbe?->?abce(pemute)
????//?success!

??????rewriter.setInsertionPointAfter(lhs.getDefiningOp());
??????auto?loc?=?NameLoc::get(rewriter.getStringAttr(lname?+?"_trans"));
??????auto?newType?=?RankedTensorType::get({lshape[0],?lshape[2],?lshape[1],?lshape[3]},?module::getElementType(lhs));
??????attrs.push_back(rewriter.getNamedAttr("order",?rewriter.getI64ArrayAttr({0,?2,?1,?3})));
??????auto?tranOp?=?rewriter.create<PermuteOp>(loc,?newType,?ValueRange{lhs},?attrs);
??????attrs.clear();
??????operands.push_back(tranOp);
??????rewriter.setInsertionPointAfter(rhs.getDefiningOp());
??????if?(auto?wOp?=?dyn_cast<top::WeightOp>(rhs.getDefiningOp()))?{

????????auto?data?=?wOp.read_as_byte();
????????uint8_t?*dptr;
????????newType?=?RankedTensorType::get({lshape[0],?rshape[0],?rshape[1],?rshape[2]},?module::getElementType(rhs));
????????std::vector<float_t>?new_filter(newType.getNumElements(),?0);
????????dptr?=?(uint8_t?*)new_filter.data();
????????for?(int32_t?i?=?0;?i?<?lshape[0];?i++)?{
??????????auto?offset?=?i?*?data->size();
??????????memcpy(dptr?+?offset,?data->data(),?data->size());
????????}
????????auto?new_op?=?top::WeightOp::create(op,?"folder",?new_filter,?newType);
????????wOp.replaceAllUsesWith(new_op.getDefiningOp());
????????operands.push_back(new_op);
????????rewriter.eraseOp(wOp);
??????}?else?{
????????loc?=?NameLoc::get(rewriter.getStringAttr(rname?+?"_reshape"));
????????newType?=?RankedTensorType::get({1,?rshape[0],?rshape[1],?rshape[2]},?module::getElementType(rhs));
????????auto?rrsop?=?rewriter.create<ReshapeOp>(loc,?newType,?ValueRange{rhs});
????????loc?=?NameLoc::get(rewriter.getStringAttr(rname?+?"_tile"));
????????attrs.push_back(rewriter.getNamedAttr("tile",?rewriter.getI64ArrayAttr({lshape[0],?1,?1,?1})));
????????newType?=?RankedTensorType::get({lshape[0],?rshape[0],?rshape[1],?rshape[2]},?module::getElementType(rhs));
????????auto?tileOp?=?rewriter.create<TileOp>(loc,?newType,?ValueRange{rrsop},?attrs);
????????attrs.clear();
????????operands.push_back(tileOp);
??????}
??????operands.push_back(none);
??????newType?=?RankedTensorType::get({lshape[0],?lshape[2],?lshape[1],?rshape[2]},?module::getElementType(op));
??????rewriter.setInsertionPoint(op);
??????loc?=?NameLoc::get(rewriter.getStringAttr(name?+?"_matmul"));
??????auto?matmulOp?=?rewriter.create<MatMulOp>(loc,?newType,?operands,?attrs);
??????attrs.clear();
??????attrs.push_back(rewriter.getNamedAttr("order",?rewriter.getI64ArrayAttr({0,?2,?1,?3})));
??????auto?tranBackOp?=?rewriter.create<PermuteOp>(op.getLoc(),?op.getType(),?ValueRange{matmulOp},?attrs);
??????op.replaceAllUsesWith(tranBackOp.getOperation());
??????rewriter.eraseOp(op);

????}?else?{
??????llvm_unreachable("Einsum?not?support?this?mode?now");
????}
????return?success();
??}

總結

TPU-MLIR對EinSum的實(shí)現雖然不完全，但已經(jīng)足夠實(shí)用，能滿(mǎn)足目前常見(jiàn)網(wǎng)絡(luò )的需求。通過(guò)Converter直接表達式規范化，降低了編譯器優(yōu)化或模式分析的復雜性。在算子分析時(shí)，我們不僅需要在計算上實(shí)現等價(jià)變換，還需充分了解實(shí)際硬件的特性。針對不同硬件架構及其對算子的支持情況，需具體分析以找到更佳實(shí)現方法。此外，我們可以看到在工程實(shí)踐中，人們更注重實(shí)用性和效率，在實(shí)現上不必追求完備，是要覆蓋實(shí)際應用場(chǎng)景即可。EinSum的轉換還有改進(jìn)空間，我們也歡迎社區提出寶貴的建議并貢獻代碼。

?

RM新时代|官方理财平台