Pytorch-mlu?實(shí)現(xiàn)添加逐層算子方法詳解
本教程分享了在寒武紀(jì)設(shè)備上 pytorch-mlu 中添加逐層算子的方法。
pytorch-mlu 逐層模式中算子間數(shù)據(jù)傳遞和存儲(chǔ)的基本單元是 tensor。pytorch-mlu 根據(jù) tensor 中的 device 屬性值將算子分發(fā)到不同設(shè)備。以 abs() 算子為例，在 dispatch 階段會(huì)根據(jù) input_tensor 的設(shè)備屬性值將算子調(diào)用分發(fā)到具體設(shè)備，邏輯如下圖所示：
Catch 通過注冊添加 MLU 算子方式與 pytorch 源碼解耦，下面介紹在 Catch 中添加 MLU 算子的具體步驟。
1、注冊算子
在 catch/torch_mlu/csrc/generated/aten_mlu_type_default.cpp 中注冊算子：
.op(torch::RegisterOperators::options().schema("aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor") // NOLINT
.impl_unboxedOnlyKernel<at::Tensor(const at::Tensor &, const at::Tensor &, at::Scalar), &AtenMluType::add>(at::TensorTypeId::MLUTensorId)
aliasAnalysis(c10::AliasAnalysisKind::FROM_SCHEMA))
2、算子分發(fā)
AtenMluType 和 AtenMluCustomType 是 Catch 模塊中算子的入口。AtenMluType 類主要包含框架中的標(biāo)準(zhǔn)算子；而 AtenMluCustomType 類包含客制化的算子。根據(jù)算子屬性選擇在 AtenMluType 還是 AtenMluCustomType 中添加相應(yīng)算子聲明和實(shí)現(xiàn)。
標(biāo)準(zhǔn)算子分發(fā)
在 catch/torch_mlu/csrc/aten/aten_mlu_type.h 和 catch/torch_mlu/csrc/aten/aten_mlu_type.cpp 中添加算子聲明和實(shí)現(xiàn)：
aten_mlu_type.h
static at::Tensor add(const at::Tensor& self, const at::Tensor& other, at::Scalar alpha);
aten_mlu_type.cpp
at::Tensor AtenMluType::add(const at::Tensor& self, const at::Tensor& other, at::Scalar alpha){
return OP_DISPATCH(add, self, other, alpha);
}
客制化算子分發(fā)
對(duì)于 MLU 特有算子，在 catch/torch_mlu/csrc/aten/aten_mlu_type.h和 catch/torch_mlu/csrc/aten/aten_mlu_custom_type.cpp 中添加算子申明和實(shí)現(xiàn)：
aten_mlu_type.h
static at::Tensor linear(const at::Tensor& input,
const at::Tensor& weight,
const at::Tensor& bias,
const at::Tensor& q_scale,
const at::Tensor& q_mode);
aten_mlu_custom_type.cpp
at::Tensor AtenMluCustomType::linear(const at::Tensor& input,
const at::Tensor& weight,
const at::Tensor& bias,
const at::Tensor& q_scale,
const at::Tensor& q_mode){
return OP_DISPATCH(linear, input, weight, bias, q_scale, q_mode);
}
3、修改 OpMethods 基類
從 AtenMluType 和 AtenMluCustomType 中都會(huì)通過 OpMethods 下發(fā)到推理算子或訓(xùn)練算子。在 catch/torch_mlu/csrc/aten/operators/op_methods.h 和 catch/torch_mlu/csrc/aten/operators/op_methods.cpp 中添加算子申明和實(shí)現(xiàn)。OpMethods 中的實(shí)現(xiàn)部分為該算子的 CPU 實(shí)現(xiàn)。
op_methods.h
virtual at::Tensor add(const at::Tensor& self, const at::Tensor& other, at::Scalar alpha);
op_methods.cpp
at::Tensor OpMethods::add(const at::Tensor& self,
const at::Tensor& other,
at::Scalar alpha){
auto input_cpu = self.cpu();
auto other_cpu = other.cpu();
auto output = at::add(input_cpu, other_cpu, alpha);
return output.to(at::Device(at::Device::Type::MLU));
}
4、下發(fā)算子
在 catch/torch_mlu/csrc/aten/operators/cnml_ops.h 和 catch/torch_mlu/csrc/aten/operators/cnml_ops.cpp 中添加推理算子申明和實(shí)現(xiàn)。
cnml_ops.h
at::Tensor add(const at::Tensor& self, const at::Tensor& other, at::Scalar alpha);
cnml_ops.cpp
at::Tensor CnmlOps::add(const at::Tensor& self, const at::Tensor& other, at::Scalar alpha){
CNML_DISPATCH(add, cnml_add, self, other, alpha); // CNML_DISPATCH 宏第一個(gè)參數(shù)是該接口名，第二個(gè)參數(shù)是wrapper個(gè)名字，其余
}
5、添加 wrapper
wrapper 是對(duì)算子 kernel 的封裝，每個(gè)算子對(duì)應(yīng)一個(gè) wrapper。這里以 add 算子為例，添加 wrapper 如下所示：
cnml_kernel.h
at::Tensor cnml_add(const at::Tensor& input, const at::Tensor& other, at::Scalar alpha);
add.cpp
at::Tensor cnml_add(const at::Tensor& input, const at::Tensor& other, at::Scalar alpha_scalar){
TORCH_CHECK(input.dim() >= 0 || other.dim() >= 0, "dimension not support");
at::Tensor input_ = input;
at::Tensor other_ = other;
auto alpha_data = alpha_scalar.to<scalar_t>();
if(alpha_data != 1){
// scale_t
other_ = cnml::ops::cnml_scale(other_, alpha_data, 0);
}
if(other_.dim() < 1 && other_.device().type() == c10::DeviceType::CPU){
auto other_scalar = other_.item();
return cnml_add_internal(input_, other_scalar); // 調(diào)用kernel
}
if(input_.dim() < 1 && input_.device().type() == c10::DeviceType::CPU){
auto input_scalar = input_.item();
return cnml_add_internal(other_, input_scalar); // 調(diào)用 kernel
}
bool broadcast = input_.sizes() != other_.sizes();
if(broadcast){
auto broadcast_size = at::infer_size(input.sizes(), other.sizes());
at::Tensor broadcast1 = cnml::ops::cnml_expand(input_, broadcast_size, false);
at::Tensor broadcast2 = cnml::ops::cnml_expand(other_, broadcast_size, false);
return cnml_add_internal(broadcast1, broadcast2); // 調(diào)用 kernel
}else{
return cnml_add_internal(input_, other_); //調(diào)用 kernel
}
return cnml_add_internal(input_, other_); //調(diào)用 kernel
}
6、添加 wrapper
Wrapper 中通過調(diào)用 kernel 實(shí)現(xiàn)算子功能。示例中調(diào)用的是 cnml_add_internal。算子的具體實(shí)現(xiàn)主要通過調(diào)用 CNML 庫的接口來完成，下面是 CNML 庫的邏輯：
kernel 實(shí)現(xiàn)就是按照上述編程邏輯調(diào)用 CNML 庫接口完成的，在 catch/torch_mlu/csrc/aten/operators/cnml/internal/cnml_internal.h 和 catch/torch_mlu/csrc/aten/operators/cnml/internal/add_internal/cpp 中添加 kernel 函數(shù)的聲明和實(shí)現(xiàn)。
cnml_internal.h
at::Tensor cnml_add_internal(const at::Tensor& input1, const at::Tensor& input2);
add_internal.cpp
at::Tensor cnml_add_internal(const at::Tensor& input1, const at::Tensor& input2){
auto output = at::native::empty_like(input1);
// prepare input cnml tensor
auto* input1_impl = getMluTensorImpl(input1); // 獲取MluTensorImpl
auto input1_cnml = input1_impl->CreateCnmlTensor(
CNML_TENSOR, toCnmlDataType(input1.dtype())); // 類型自適應(yīng)：toCnmlDataType()
auto* input2_impl = getMluTensorImpl(input2);
auto input2_cnml = input2_impl->CreateCnmlTensor(
CNML_TENSOR, toCnmlDataType(input2.dtype()));
// prepare output cnml tensor
auto* output_impl = getMluTensorImpl(output);
auto output_cnml = output_impl->CreateCnmlTensor(
CNML_TENSOR, toCnmlDataType(output.dtype()));
// End the execution flow if not MLU device
CHECK_MLU_DEVICE(output);
// setup operator
cnmlBaseOp_t add_op;
TORCH_CNML_CHECK(cnmlCreateAddOp(&add_op, input1_cnml, input2_cnml, output_cnml));
// return to JIT if running mode is fuse
CHEXK_RETURN_TO_FUSE(add_op, output);
// compile op
TORCH_CNML_CHECK(cnmlCompileBaseOp(add_op, GET_CORE_VERSION, GET_CORE_NUMBER));
auto queue = getCurQueue();
TORCH_CNML_CHECK(cnmlComputeAddOpForward_V4(add_op,
NULL,
input1_impl->raw_mutable_data(),
NULL,
input2_impl->raw_mutable_data(),
NULL,
output_impl->raw_mutable_data(),
queue,
NULL));
syncQueue(queue);
TORCH_CNML_CHECK(cnmlDestroyBaseOp(&add_op));
return output;
}
對(duì) MLU 不支持算子的處理
對(duì)于 MLU 暫不支持的操作，輸入數(shù)據(jù)將會(huì)拷貝到 CPU 上，然后調(diào)用 CPU 相關(guān)操作，使其在 CPU 上運(yùn)行，最后再將輸出結(jié)果拷會(huì)到 MLU 上。具體實(shí)現(xiàn)，可以查詢 op_methods.cp，該文件在 catch/torch_mlu/csrc/aten/operators/ 目錄下。
op_methods.cpp
at::Tensor OpMethods::add(const at::Tensor& self,
const at::Tensor& other,
at::Scalar alpha){
auto input_cpu = self.cpu();
auto other_cpu = other.cpu();
auto output = at::add(input_cpu, other_cpu, alpha);
return output.to(at::Device(at::Device::Type::MLU));
}
對(duì)于新增的算子在執(zhí)行過程中拋出異常時(shí)，如果 CPU 上沒有對(duì)應(yīng)的算子操作，那么該操作無法切換到 CPU 上運(yùn)行；
Wrapper一般以 cnml_算子名命名，kernel一般以cnml_算子名_internal命名
7、算子測試
使用基于 python 的 unittest 模塊編寫算子單元測試。測試時(shí)需提供相同的參數(shù)和輸入數(shù)據(jù)，分別在 MLU 和 CPU 上執(zhí)行算子，對(duì)比兩者的輸出結(jié)果。MLU 和 CPU 計(jì)算結(jié)果可能會(huì)有差異，一般情況下兩者的相對(duì)誤差在 2% 以內(nèi)均是可以接受的。
def test_add(self):
# "Tensor + Tensor" mode testing
for shape1, shape2 in [((1,3,224,224),(1,3,224,224)),((2,30,80),(2,30,80)),((3,20),(3,20)),((10),(10))]:
input1_cpu = torch.rand(shape1, dtype=torch.float)
input2_cpu = torch.rand(shape2, dtype=torch.float)
input1_mlu = input1_cpu.to(xm.mlu_device())
input2_mlu = input2_cpu.to(xm.mlu_device())
# 在 CPU 上計(jì)算
output_cpu = input1_cpu + input2_cpu
# 在 MLU 上計(jì)算
output_mlu = input1_mlu + input2_mlu
# 計(jì)算 MLU 的誤差，并確保相對(duì)誤差在 2% 以內(nèi)
self.assertTensorsEqual(output_cpu, output_mlu.cpu(), 0.02, use_MSE=True)
?以上分享了在寒武紀(jì)設(shè)備 pytorch-mlu 中添加逐層算子的方法，并以 add() 算子為例進(jìn)行了示例編寫，希望我的分享會(huì)對(duì)你的學(xué)習(xí)有一點(diǎn)幫助。
到此這篇關(guān)于Pytorch-mlu 實(shí)現(xiàn)添加逐層算子方法詳解的文章就介紹到這了,更多相關(guān)Pytorch內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！
相關(guān)文章
這篇文章主要介紹了介紹Python的@property裝飾器的用法,是Python學(xué)習(xí)進(jìn)階中的重要知識(shí),代碼基于Python2.x版本,需要的朋友可以參考下2015-04-04
python 動(dòng)態(tài)遷移solr數(shù)據(jù)過程解析
這篇文章主要介紹了python 動(dòng)態(tài)遷移solr數(shù)據(jù)過程解析,文中通過示例代碼介紹的非常詳細(xì)，對(duì)大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值,需要的朋友可以參考下2019-09-09
使用Windows批處理和WMI設(shè)置Python的環(huán)境變量方法
今天小編就為大家分享一篇使用Windows批處理和WMI設(shè)置Python的環(huán)境變量方法，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。一起跟隨小編過來看看吧2019-08-08
windows下 兼容Python2和Python3的解決方法
這篇文章主要介紹了windows下 兼容Python2和Python3的解決方法,需要的朋友可以參考下2018-12-12
利用python查找電腦里的文件非常方便，首先是導(dǎo)入模塊os，然后利用os.listdir()打開目錄下的所有文件，files其實(shí)是一個(gè)列表。具體操作方法大家參考下本文2018-04-04
這篇文章主要帶大家來復(fù)習(xí)一下Python基礎(chǔ)中的列表操作，不知道各位還記得多少呢？文中的示例代碼講解詳細(xì)，對(duì)我們學(xué)習(xí)Python有一定幫助，需要的可以參考一下2022-11-11
這篇文章主要介紹了Python之Matlibplot畫圖功能演示過程,具有很好的參考價(jià)值,希望對(duì)大家有所幫助,如有錯(cuò)誤或未考慮完全的地方,望不吝賜教2023-09-09
python實(shí)現(xiàn)會(huì)員信息管理系統(tǒng)(List)
這篇文章主要為大家詳細(xì)介紹了python實(shí)現(xiàn)會(huì)員信息管理系統(tǒng)，文中示例代碼介紹的非常詳細(xì)，具有一定的參考價(jià)值，感興趣的小伙伴們可以參考一下2022-03-03
Pytorch出現(xiàn)錯(cuò)誤Attribute?Error:module?‘torch‘?has?no?attrib
這篇文章主要給大家介紹了關(guān)于Pytorch出現(xiàn)錯(cuò)誤Attribute?Error:module?‘torch‘?has?no?attribute?'_six'解決的相關(guān)資料,文中通過圖文介紹的非常詳細(xì),需要的朋友可以參考下2023-11-11
這篇文章主要介紹了python處理xml文件的方法,結(jié)合實(shí)例形式總結(jié)分析了Python常見的xml文件處理技巧與相關(guān)注意事項(xiàng),需要的朋友可以參考下2017-05-05