python教程分享Pytorchmlu 实现添加逐层算子方法详解

目录

• 1、注册算子
• 2、算子分发
• 3、修改 opmethods 基类
• 4、下发算子
• 5、添加 wrapper
• 6、添加 wrapper
• 7、算子测试

  .op(torch::registeroperators::options().schema("aten::add.tensor(tensor self, tensor other, *, scalar alpha=1) -> tensor")  // nolint       .impl_unboxedonlykernel(at::tensortypeid::mlutensorid)        aliasanalysis(c10::aliasanalysiskind::from_schema))  

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

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

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

  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 宏第一个参数是该接口名，第二个参数是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();    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);   // 调用kernel    }    if(input_.dim() <1 && input_.device().type() == c10::devicetype::cpu){      auto input_scalar = input_.item();      return cnml_add_internal(other_, input_scalar);   // 调用 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);  // 调用 kernel    }else{      return cnml_add_internal(input_, other_);  //调用 kernel    }    return cnml_add_internal(input_, other_);   //调用 kernel  }  

  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()));  // 类型自适应：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;  }  

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

  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 上计算      output_cpu = input1_cpu + input2_cpu      # 在 mlu 上计算      output_mlu = input1_mlu + input2_mlu      # 计算 mlu 的误差，并确保相对误差在 2% 以内      self.asserttensorsequal(output_cpu, output_mlu.cpu(), 0.02, use_mse=true)