[Обмен опытом] Cambrian pytorch-mlu добавляет послойный операторный метод

PyTorch
[Обмен опытом] Cambrian pytorch-mlu добавляет послойный операторный метод

 欢迎关注我的公众号 [极智视界]，回复001获取Google编程规范

 O_o >_<  o_O O_o ~_~ o_O

  В этом руководстве рассказывается о методе добавления послойных операторов в pytorch-mlu на устройствах Cambrian.

  Базовой единицей передачи и хранения данных между операторами в послойном режиме pytorch-mlu является тензор. pytorch-mlu распределяет операторов по разным устройствам в соответствии со значением атрибута устройства в тензоре. Взяв в качестве примера оператор abs(), на этапе отправки вызов оператора будет распределяться по конкретным устройствам в соответствии со значением атрибута устройства input_tensor, Логика показана на следующем рисунке:

  Catch отделяется от исходного кода pytorch путем регистрации и добавления операторов MLU. Ниже описаны конкретные шаги по добавлению операторов MLU в Catch.

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. Распределение оператора

  AtenMluType и AtenMluCustomType — это записи операторов в модуле Catch. Класс AtenMluType в основном содержит стандартные операторы платформы, а класс AtenMluCustomType содержит настраиваемые операторы. Выберите, следует ли добавить соответствующее объявление оператора и реализацию в AtenMluType или AtenMluCustomType в соответствии со свойствами оператора.

  • Стандартное операторское распределение существуетcatch/torch_mlu/csrc/aten/aten_mlu_type.hиcatch/torch_mlu/csrc/aten/aten_mlu_type.cppДобавьте объявление оператора и реализацию в:
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);
}
  • Индивидуальное распределение операторов

  Для операторов, специфичных для MLU, вcatch/torch_mlu/csrc/aten/aten_mlu_type.hиcatch/torch_mlu/csrc/aten/aten_mlu_custom_type.cppДобавьте объявление оператора и реализацию в:

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 отправляются оператору логического вывода или оператору обучения через OpMethods. существуетcatch/torch_mlu/csrc/aten/operators/op_methods.hиcatch/torch_mlu/csrc/aten/operators/op_methods.cppДобавьте объявление оператора и реализацию в . Часть реализации в OpMethods — это реализация оператора процессором.

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. Выдать оператор

существуетcatch/torch_mlu/csrc/aten/operators/cnml_ops.hиcatch/torch_mlu/csrc/aten/operators/cnml_ops.cppДобавлено объявление и реализация оператора вывода в .

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个名字，其余
}

5. Добавьте обертку

  обертка — это инкапсуляция ядра оператора, каждому оператору соответствует обертка. Здесь мы возьмем оператор добавления в качестве примера, добавление обертки выглядит следующим образом:

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);   // 调用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
}

6. Добавьте обертку

  Оболочка реализует операторные функции, вызывая ядро. В примере вызывается cnml_add_internal. Конкретная реализация оператора в основном завершается вызовом интерфейса библиотеки CNML.Следующая логика библиотеки CNML:

  Реализация ядра осуществляется путем вызова интерфейса библиотеки CNML в соответствии с приведенной выше логикой программирования.catch/torch_mlu/csrc/aten/operators/cnml/internal/cnml_internal.hиcatch/torch_mlu/csrc/aten/operators/cnml/internal/add_internal/cppДобавьте объявление и реализацию функции ядра в .

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;
}
  • Обработка операторов, не поддерживаемых MLU

  Для операций, которые в настоящее время не поддерживаются MLU, входные данные будут скопированы в ЦП, а затем будут вызваны операции, связанные с ЦП, для выполнения на ЦП, и, наконец, выходные результаты будут скопированы в MLU. Для конкретной реализации вы можете запросить 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));
}
  • Когда возникает исключение во время выполнения вновь добавленного оператора, если на ЦП нет соответствующей операции оператора, операция не может быть переключена на выполнение на ЦП;
  • Обертка обычно начинается сcnml_Оператор назван в честь ядра в целомcnml_имя оператора_internalназвание

7. Тест оператора

  Пишите юнит-тесты операторов, используя модуль unittest на основе Python. Во время теста необходимо предоставить одни и те же параметры и входные данные, операторы выполняются на MLU и ЦП соответственно, а выходные результаты двух сравниваются. Результаты расчета MLU и CPU могут отличаться, как правило, относительная погрешность между ними допустима в пределах 2%.

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)

  Выше был описан метод добавления послойных операторов в кембрийском устройстве pytorch-mlu, и в качестве примера был написан пример с оператором add(). Надеюсь, мой обмен немного поможет вам в вашем обучении.

【Передача по общему номеру】 "[Обмен опытом] Cambrian pytorch-mlu добавляет послойный операторный метод》

Категории