Examples

// add_custom_kernel/custom_op.cpp
// Shape inference function of the AddCustom operator, which implements the PyTorch-style broadcast semantics. GE calls this function during graph compilation to infer the output shape based on the input shape.
graphStatus InferShapeForAdd(gert::InferShapeContext *context) {
  // Obtain the shapes of the two inputs.
  auto input_shape_0 = *context->GetInputShape(0);
  auto input_shape_1 = *context->GetInputShape(1);
  // Obtain the writable reference of the output shape, which is used to set the inference result.
  auto output_shape = context->GetOutputShape(0);
  // Case 1: The two inputs have different dimensions. Handle the broadcast semantics (broadcasting a 1D tensor to a multi-dimensional tensor).
  if (input_shape_0.GetDimNum() != input_shape_1.GetDimNum()) {
    auto min_num = std::min(input_shape_0.GetDimNum(), input_shape_1.GetDimNum());
    auto max_num = std::max(input_shape_0.GetDimNum(), input_shape_1.GetDimNum());
    if (min_num != 1) {
      std::cerr << "Add param invalid" << std::endl;
    } else {
      if (input_shape_1.GetDimNum() > 1) {
        *output_shape = input_shape_1;
      } else {
        *output_shape = input_shape_0;
      }
      return GRAPH_SUCCESS;
    }
  }
  // Case 2: If one input is a scalar (dimension count = 0), the output shape is the shape of the other input.
  if (input_shape_0.GetDimNum() == 0) {
    *output_shape = input_shape_1;
    return GRAPH_SUCCESS;
  }
  // Case 3: If the two inputs have the same number of dimensions, the maximum value is taken dimension-wise (standard broadcasting rule).
  if (input_shape_1.GetDimNum() == 0) {
    *output_shape = input_shape_0;
    return GRAPH_SUCCESS;
  }
  output_shape->SetDimNum(input_shape_0.GetDimNum());
  for (size_t i = 0; i < input_shape_0.GetDimNum(); ++i) {
    output_shape->SetDim(i, std::max(input_shape_0.GetDim(i), input_shape_1.GetDim(i)));
  }
  return GRAPH_SUCCESS;
}

// Register the InferShape function with the AddCustom operator. When encountering an AddCustom node during GE compilation, the InferShape function is called to infer the output shape.
IMPL_OP(AddCustom).InferShape(InferShapeForAdd);

// add_custom_kernel/add_custom.asc
// Operator implementation on the PyTorch: Execute the Ascend C kernel on the NPU.
namespace ascendc_ops {
// ascendc_add function: called by PyTorch in eager mode.
// Prepare parameters and start the Ascend C kernel function.
at::Tensor ascendc_add(const at::Tensor& x, const at::Tensor& y) {
  // Obtain the current NPU execution stream. The kernel is executed asynchronously on this stream.
  auto aclStream = c10_npu::getCurrentNPUStream().stream(false);
  // Allocate the output tensor, with the shape the same as that of the input x.
  at::Tensor z = at::empty_like(x);
  // Set the number of blocks to be started by the kernel function.
  uint32_t numBlocks = 8;
  // Calculate the total number of input elements (by multiplying the number of all dimensions).
  uint32_t totalLength = 1;
  for (uint32_t size : x.sizes()) {
    totalLength *= size;
  }
  // Call the encapsulation function to start the kernel function.
  launch_add_custom(
    static_cast<uint8_t*>(x.data_ptr()),  // Device address of input x
    static_cast<uint8_t*>(y.data_ptr()),  // Device address of input y
    static_cast<uint8_t*>(z.data_ptr()),  // Device address of output z
    totalLength,                          // Total number of elements
    numBlocks,                            // Number of blocks of the kernel function
    aclStream                             // Execution stream
  );
  return z;
}
} // namespace ascendc_ops

// ===== PyTorch operator registration =====
// Define the operator schema: Declare the function signature for PyTorch to identify.
// ascendc_ops is the namespace, and ascendc_add is the function name.
TORCH_LIBRARY_FRAGMENT(ascendc_ops, m) {
  m.def("ascendc_add(Tensor x, Tensor y) -> Tensor");
}

// Register the implementation with PrivateUse1: eager execution path on the NPU device.
// Triggered when a user directly calls torch.ops.ascendc_ops.ascendc_add(x, y) on the NPU.
TORCH_LIBRARY_IMPL(ascendc_ops, PrivateUse1, m) {
  m.impl("ascendc_add", TORCH_FN(ascendc_ops::ascendc_add));
}

// Register the implementation with XLA: graph compilation path.
// When the operator enters the graph mode through torch.compile, the XLA path triggers GE graph compilation.
TORCH_LIBRARY_IMPL(ascendc_ops, XLA, m) {
  m.impl("ascendc_add", TORCH_FN(ascendc_ops::ascendc_add));
}

// Register the implementation for Meta: shape inference during compile.
// The Meta path does not perform actual computation and only returns an empty tensor with the same shape as the input.
// TorchAir calls this implementation for shape inference during graph compilation.
TORCH_LIBRARY_IMPL(ascendc_ops, Meta, m) {
  m.impl("ascendc_add", [](const torch::Tensor& x, const torch::Tensor& y) {
    return at::empty_like(x);
  });
}