Other Mapping Information

Callback APIs

The following describes the implementation details and related APIs of each callback API.

EagerExecuteOp::Execute

This is the core API and is called back during the runtime. It must be implemented in all scenarios.

graphStatus Execute(gert::EagerOpExecutionContext *ctx) override;

Context APIs:

Method	Function
ctx->GetInputTensor(index)	Obtains the input tensor (including the addr, shape, dtype, and format).
ctx->MallocOutputTensor(index, shape, format, dtype)	Allocates the device memory for the output tensor.
ctx->MallocWorkSpace(size)	Allocates the workspace device memory.
ctx->GetStream()	Obtains the stream.
ctx->GetOutputTensor(index)	Obtains the allocated output tensor.
ctx->MakeOutputRefInput(out_idx, in_idx)	Outputs the reference input address (no copy required).

Typical process:

graphStatus Execute(gert::EagerOpExecutionContext *ctx) {
  // Obtain the 0th input tensor (a constant that cannot be modified).
  auto *x = ctx->GetInputTensor(0);
  // Allocate the output tensor memory. The shape, format, and dtype are inherited from the input.
  auto *y = ctx->MallocOutputTensor(0, x->GetShape(), x->GetFormat(), x->GetDataType());
  // Obtain the stream where the kernel is executed asynchronously.
  void *stream = ctx->GetStream();
  // Start the custom kernel.
  LaunchMyKernel(x->GetAddr(), y->GetAddr(), stream);
  return GRAPH_SUCCESS;
}

Precautions:

The output memory is managed by the context, and does not need to be released by the caller.
GetInputTensor returns const Tensor*, and the input data cannot be modified.
In the precompiled kernel binary scenario, load the binary in Execute and launch it using aclrtLaunchKernelWithHostArgs.

ShapeInferOp::InferShape/InferDataType

This callback API is used for shape and data type inference during compilation.

graphStatus InferShape(gert::InferShapeContext *ctx) override;
graphStatus InferDataType(gert::InferDataTypeContext *ctx) override;

InferShapeContext API:

Method	Function
ctx->GetInputShape(index)	Obtains the input origin shape (const Shape*).
ctx->GetOutputShape(index)	Obtains the modifiable output shape (Shape*).
ctx->GetInputTensor(index)	Obtains the input tensor (accessible data).

InferDataTypeContext API:

Method	Function
ctx->GetInputDataType(index)	Obtains the input data type.
ctx->SetOutputDataType(index, dtype)	Sets the output data type.

Typical process:

// Shape inference: Set the output shape to be the same as the input shape.
graphStatus InferShape(gert::InferShapeContext *ctx) {
  auto *in_shape = ctx->GetInputShape(0);
  auto *out_shape = ctx->GetOutputShape(0);
  *out_shape = *in_shape;
  return GRAPH_SUCCESS;
}

// DataType inference: Set the output data type to be the same as the input data type.
graphStatus InferDataType(gert::InferDataTypeContext *ctx) {
  return ctx->SetOutputDataType(0, ctx->GetInputDataType(0));
}

CompilableOp::Compile

This is an online compilation callback API, which is called during GE/ATC compilation.

graphStatus Compile(gert::OpCompileContext *ctx) override;

OpCompileContext API:

Method	Function
ctx->GetInputTensor(index)	Obtains the input tensor (including the shape, dtype, and format) during compilation.
ctx->GetOutputTensor(index)	Obtains the output tensor during compilation.
ctx->GetOption(key, value)	Obtains compilation options.
ctx->GetPlatformInfos(...)	Obtains the target platform information (such as the chip model).

Typical process:

graphStatus Compile(gert::OpCompileContext *ctx) {
  // Obtain the input tensor (including the shape, dtype, and format) during compilation.
  auto *input = ctx->GetInputTensor(0);
  // Generate the binary cache key based on the input shape to support kernel management in multi-shape scenarios.
  auto key = BuildBinaryKey(input->GetShape());
 // Load the kernel source code file.
  auto source = LoadFile("my_kernel.cpp");

  // Compile the kernel source code into a device binary using Runtime Compilation (RTC).
  aclrtcProg prog;
  aclrtcCreateProg(&prog, source.c_str(), "kernel_name", 0, nullptr, nullptr);
  aclrtcCompileProg(prog, num_options, options);

  // Obtain the compilation artifact (device binary).
  size_t bin_size;
  aclrtcGetBinDataSize(prog, &bin_size);
  std::vector<uint8_t> binary(bin_size);
  aclrtcGetBinData(prog, binary.data());
 // Cache the compilation artifact to a member variable for subsequent execution.
  device_elves_[key] = std::move(binary);
  aclrtcDestroyProg(&prog);   // Destroy the compilation program object.
  return GRAPH_SUCCESS;
}

PortableOp::Serialize / Deserialize

This is a serialization and deserialization callback API, which is used for OM model offloading.

// Serialization: Write the compilation artifact to the buffer for ATC to write to an OM model file.
// The buffer format is user-defined. GE only passes the buffer through without parsing it.
graphStatus Serialize(std::vector<uint8_t> &buffer) override {
  // Serialization format: [count][key_len][key][bin_len][bin]...
  size_t offset = 0;
  buffer.resize(CalcTotalSize());
  
  // Write the number of binaries.
  uint32_t count = device_elves_.size();
  memcpy(buffer.data() + offset, &count, sizeof(count));
  offset += sizeof(count);
  
  // Write the key and data for each binary one by one.
  for (auto &kv : device_elves_) {
    //... Write the keys and binaries.
  }
  return GRAPH_SUCCESS;
}

// Deserialization: Restore the compilation artifact from the buffer of the OM model file.
// The format must exactly match the counterpart in serialization.
graphStatus Deserialize(const std::vector<uint8_t> &buffer) override {
  size_t offset = 0;
  
  // Read the number of binaries.
  uint32_t count;
  memcpy(&count, buffer.data() + offset, sizeof(count));
  offset += sizeof(count);
  
  // Read the key and data for each binary one by one.
  for (uint32_t i = 0; i < count; ++i) {
    //... Read the keys and binaries and store them in device_elves_.
  }
  return GRAPH_SUCCESS;
}

Key points:

The buffer format is user-defined. GE only passes the buffer through without parsing it.
Serialize is called in the ATC phase to write the compilation artifact to the buffer.
Deserialize is called in the OM loading and execution phase to restore the compilation artifact from the buffer.
Recommended serialization format in multi-binary scenarios: [count][key_len][key][bin_len][bin]...

ArgsUpdater::UpdateHostArgs

This API is called back to update the address reference in kernel arguments when an I/O address changes.

graphStatus UpdateHostArgs(gert::UpdateArgsContext *ctx) override;

Context APIs:

Method	Function
ctx->GetKernelArgs(placement, index)	Obtains the kernel arguments at a specified position.
All methods inherited from EagerOpExecutionContext	Obtains information such as the input and output tensors.

Application scenario: After static graph offloading, the I/O address during model loading may be different from that during compilation. The address in the args buffer needs to be updated.

Debugging and Common Errors

Development Check Items
Check by graph integration phase:
- Phases 1 and 2 (definition and bridging):
  - The operator type name, registration class name, and operator type used during graph construction must be the same.
  - PyTorch scenario: The TORCH_LIBRARY schema must be consistent with the input and output in REG_OP.
  - The GE operator type in the converter must be the same as the registration name in REG_AUTO_MAPPING_OP.
- Phase 3 (graph integration):
  - PyTorch: The converter must be registered, and torchair.ge_mode() must be enabled.
- Phases 4 and 5 (compilation and execution):
  - ASCEND_CUSTOM_OPP_PATH must point to the root directory of the OPP package.
  - op_graph/lib/<os>/<arch>/ must contain the .so file.
  - The kernel binary/source code path must not depend on temporary working directories.
  - In the offline OM scenario, ensure that the value of soc_version is consistent with that in the actual environment.

Common errors

Symptom	Possible Cause	Troubleshooting
The GE cannot find any custom operator.	The ASCEND_CUSTOM_OPP_PATH environment variable is not configured, or the path is incorrect.	Check the environment variable to ensure that the op_graph/lib/<os>/<arch>/ contains the .so file.
Execute returns null input.	The input index does not match the REG_OP definition.	Check the INPUT sequence of REG_OP and the GetInputTensor index in Execute.
The shape is empty during compilation.	The input shape is not available during compilation.	Check whether GetInputTensor is correctly read during compilation.
The OM failed to be executed after being loaded.	Deserialization does not correctly restore binaries.	Check whether the buffer formats of serialization and deserialization are symmetric.
A PyTorch graph mode error is reported.	The converter is not registered, or the operator type name does not match.	Ensure that the GE operator type in register_fx_node_ge_converter is the same as that in REG_AUTO_MAPPING_OP.

Parent topic: Integrating Custom Operators into a Graph for PyTorch