CustomOp API Examples

OutputMetaFn

The following examples demonstrate callback function returning a vector of output tensors parameters.

• Example for op returning two output tensors:

  static habana::PartialOutputMetaDataVector output_meta(
      const at::Stack& inputs) {
  
    // Acquire inputs needed to calculate outputs parameters.
    auto self = inputs[0].toTensor();
    auto k = inputs[1].toInt();
    auto dim = inputs[2].toInt();
  
    // Calculate output shape.
    std::vector<int64_t> output_shape = self.sizes().vec();
    if (output_shape.size() > 0) {
      output_shape[dim] = k;
    }
  
    // Define metadata structures containing dtypes and shapes of output tensors.
    habana::PartialOutputMetaData meta_output{
        c10::ScalarType::Float, output_shape};
    habana::PartialOutputMetaData meta_indices{
        c10::ScalarType::Long, output_shape};
  
    // Return these structures as a vector.
    return {meta_output, meta_indices};
  }
  

• Example for op returning one output tensor:

  static habana::PartialOutputMetaDataVector output_meta(
      const at::Stack& inputs) {
  
    // Acquire inputs needed to calculate outputs parameters.
    auto self = inputs[0].toTensor();
    auto other = inputs[1].toTensor();
  
    // Calculate output shape.
    auto output_shape = at::infer_size(self.sizes(), other.sizes());
  
    // Define metadata structure containing dtype and shape of the output tensor.
    habana::PartialOutputMetaData meta_output{self.scalar_type(), output_shape};
  
    // Return it as a vector.
    return {meta_output};
  }
  

FillParamsFn

The following examples demonstrate callback function returning a kernel params structure.

• Example for op calling TPC kernel:

  #include <perf_lib_layer_params.h>
  
  static std::shared_ptr<void> fill_params(
      const at::Stack& inputs,
      size_t& size) {
  
    // Create a params structure using helper macro.
    HPU_PARAMS_STUB(ns_ReluKernel::Params);
  
    // Fill params fields.
    params->threshold.f = 0.0;
  
    // Return the structure.
    return params;
  }
  

• Example for the op calling synapse kernel:

  #include <synapse_common_types.hpp>
  
  static std::shared_ptr<void> fill_params(
      const at::Stack& inputs,
      size_t& size) {
  
    // Create a params structure using helper macro.
    HPU_PARAMS_STUB(synBeamParams);
  
    auto self = inputs[0].toTensor();
  
    // Fill params fields.
    params->bsw = inputs[1].toInt();
    auto dim = inputs[2].toInt();
    params->axis = self.dim() - dim - 1;
    params->bottomK = inputs[3].toBool();
  
    // Return the structure.
    return params;
  }
  

Meta Implementation

The following examples demonstrate function defining output tensors for torch.compile.

• Example for op returning two output tensors:

  std::tuple<at::Tensor, at::Tensor> custom_topk_meta(
      at::Tensor input_a,
      at::Scalar k,
      at::Scalar axis,
      bool bottom) {
  
    // Calculate output shape.
    auto output_shape = input_a.sizes().vec();
    if (output_shape.size() > 0) {
      output_shape[axis.toInt()] = k.toInt();
    }
  
    // Create empty at::Tensors with the shapes and dtypes reflecting
    // operator's output tensors.
    auto output = input_a.new_empty(output_shape, c10::ScalarType::Float);
    auto indices = input_a.new_empty(output_shape, c10::ScalarType::Long);
  
    // Return them as a vector.
    return {output, indices};
  }
  
  /**
  * Register defined function for the Meta DispatchKey.
  */
  TORCH_LIBRARY_IMPL(custom_op, Meta, m) {
    m.impl("custom_topk", &custom_topk_meta);
  }
  

• Example for op returning one output tensor:

  at::Tensor custom_add_meta(at::Tensor input_a, at::Tensor input_b) {
      // Calculate output shape.
  
    auto output_shape = at::infer_size(input_a.sizes(), input_b.sizes());
  
    // Return empty at::Tensor with the shape and dtype reflecting
    // operator's output tensor.
    return input_a.new_empty(output_shape, input_a.scalar_type());
  }
  
  
  /**
  * Register defined function for the Meta DispatchKey.
  */
  TORCH_LIBRARY_IMPL(custom_op, Meta, m) {
      m.impl("custom_add", &custom_add_meta);
  }