Graph Operator
Procedure
- The following code is a simple example. The return value of the API function is not verified. The return value needs to be verified in actual use.
- When you use the ATB for development, use tensors as the inputs.
- Include the ACL and ATB API header files.
#include <acl/acl.h> #include <atb/atb_infer.h>
- Set deviceId.Set deviceId as required.
int deviceId = 0; aclError status = aclrtSetDevice(deviceId);
- Create a context and configure the stream.The context manages the streams used by the NPU.
atb::Context *context = nullptr; st = atb::CreateContext(&context); aclrtStream stream = nullptr; status = aclrtCreateStream(&stream); context->SetExecuteStream(stream);
- Create a GraphOperation object instance.
You can create and use a graph operator by configuring TensorId or TensorName. For details about the graph operator structure, see Figure 1. For details about the mapping between TensorId and TensorName, see Table 1.
Table 1 Mapping between tensor IDs and names Tensor
TensorId
TensorName
a
0
"a"
b
1
"b"
c
2
"c"
output
3
"output"
a_add_b_output
4
"a_add_b_output"
- Mode 1: Configure TensorId.
- Construct an Operation parameter.
Different from the parameters of a single-operator, the parameters of a graph operator include graph-related information such as the graph node, number of input tensors, number of output tensors, and number of intermediate tensors.
Calculate the number of input tensors (x), output tensors (y), and intermediate tensors (z) based on the designed graph operator structure. The ID of the input tensor is in the range of [0, x - 1], the ID of the output tensor is in the range of [x, x + y - 1], and the ID of the intermediate tensor is in the range of [x + y, x + y + z - 1]. For details about the mapping, see Tensor and TensorId columns in Table 1.
Then, configure the information about each node, including the created single-operator object instance, input tensor, and output tensor. The input and output tensors of this node may be the input tensor, output tensor, or intermediate tensor of the graph. You need to select a value within a proper range based on the graph tensor type.
For details about how to create op0 and op1 in the instance, see the creation of a single-operator.
atb::GraphParam graphParam; graphParam.inTensorNum = 3; // Specify the number of input tensors of the graph. graphParam.outTensorNum = 1; // Specify the number of output tensors of the graph. graphParam.internalTensorNum = 1; // Specify the number of intermediate tensors of the graph. graphParam.nodes.resize(2); // Specify the number of nodes in the graph, that is, the number of single-operators. graphParam.nodes[0].operation = op0; // Specify the single-operator object instance of node 0 in the graph. graphParam.nodes[0].inTensorIds = {0, 1}; // Specify the ID of the input tensor required by node 0 in the graph. graphParam.nodes[0].outTensorIds = {4}; // Specify the ID of the output tensor of node 0 in the graph. graphParam.nodes[1].operation = op1; // Specify the single-operator object instance of node 1 in the graph. graphParam.nodes[1].inTensorIds = {4, 2}; // Specify the ID of the input tensor required by node 1 in the graph. graphParam.nodes[1].outTensorIds = {3}; // Specify the ID of the output tensor of node 1 in the graph. - Create an operator object instance.
atb::Operation *op = nullptr; atb::Status st = atb::CreateOperation(graphParam, &op);
- Construct an Operation parameter.
- Mode 2: Configure TensorName.
TensorId needs to be defined in advance, and the operation process is complex. This mode defines each tensor by using a character string, which is more feasible. For details about the mapping, see Tensor and TensorName columns in Table 1.
- Create a graph operator builder.
atb::GraphOpBuilder* graphOpBuilder; CreateGraphOpBuilder(&graphOpBuilder);
- Initialize the graph operator builder.
// lambda function, which is used to deduce the output tensorDesc from the input tensorDesc of the graph operator, including DataType, Format, and Shape. atb::InferShapeFunc inferShapeFunc = [=](const atb::SVector<atb::TensorDesc> &inTensorDescs, atb::SVector<atb::TensorDesc> &outTensorDescs) { outTensorDescs.at(0) = inTensorDescs.at(0); return atb::NO_ERROR; }; graphOpBuilder->Init("DemoGraphOperation", inferShapeFunc, {"a", "b", "c"}, {"output"}); - Use the graph operator constructor to construct a graph.
During graph construction, you can define the lambda function to reshape the tensor. Ensure that the shape size before and after reshaping is the same.
For details about how to create a single-operator such as op0, see the creation of a single-operator.
graphOpBuilder->AddOperation(op0, {"a", "b"}, {"a_add_b_output"}); graphOpBuilder->AddOperation(op1, {"a_add_b_output", "c"}, {"output"}); - Create the instance of a graph operator object.
atb::Operation *op = graphOpBuilder->Build(); // When using it, you need to determine whether op is a null pointer. DestroyGraphOpBuilder(graphOpBuilder); // Destroy the graph operator builder.
- Create a graph operator builder.
- Mode 1: Configure TensorId.
- Create a VariantPack.VariantPack contains the input and output tensor lists. Each input tensor passed in VariantPack must be greater than 0 and less than or equal to 256 GB.
// Tensor construction method atb::Tensor a; a.desc.dtype = ACL_FLOAT16; // Configure the tensor data type. a.desc.format = ACL_FORMAT_ND; // Configure the tensor format. a.desc.shape.dimNum = 2; // Configure the number of tensor dimensions. a.desc.shape.dims[0] = 3; // Configure the size of dimension 0 of the tensor. a.desc.shape.dims[1] = 3; // Configure the size of the first dimension of the tensor. a.dataSize = Utils::GetTensorSize(a); // Obtain the memory size of a tensor. status = aclrtMalloc(&a.deviceData, a.dataSize, ACL_MEM_MALLOC_HUGE_FIRST); // Allocate the device memory. // Construct all input and output tensors based on the preceding method and save them to VariantPack. atb::VariantPack variantPack; atb::SVector<atb::TensorDesc> intensorDescs; atb::SVector<atb::TensorDesc> outtensorDescs; uint32_t inTensorNum = operation->GetInputNum(); uint32_t outTensorNum = operation->GetOutputNum(); pack.inTensors.resize(inTensorNum); intensorDescs.resize(inTensorNum); // Create the description of each input tensor. Each tensor needs to match the data type (dtype), data format (format), and data dimensions (dimNum and dims). intensorDescs.at(0).dtype = ACL_FLOAT16; intensorDescs.at(0).format = ACL_FORMAT_ND; intensorDescs.at(0).shape.dimNum = 2; intensorDescs.at(0).shape.dims[0] = 2; intensorDescs.at(0).shape.dims[1] = 2; inTensors.at(0).desc = intensorDescs.at(0); inTensors.at(0).dataSize = atb::Utils::GetTensorSize(inTensors.at(0)); // Copy the created tensor to variantPack. aclrtMemcpy( pack.inTensors.at(0).deviceData, pack.inTensors.at(i).dataSize, a.data(), a.size(), ACL_MEMCPY_HOST_TO_DEVICE); // Place the description and data size of other input tensors. outtensorDescs.resize(outTensorNum); pack.outTensors.resize(outTensorNum); operation->InferShape(intensorDescs, outtensorDescs); // Create the description of each output tensor. Each tensor needs to match the data type (dtype), data format (format), and data dimensions (dimNum and dims). outtensorDescs.at(0).dtype = ACL_FLOAT16; outtensorDescs.at(0).format = ACL_FORMAT_ND; outtensorDescs.at(0).shape.dimNum = 2; outtensorDescs.at(0).shape.dims[0] = 3; outtensorDescs.at(0).shape.dims[1] = 4; pack.outTensors.at(0).desc = outtensorDescs.at(0); pack.outTensors.at(0).dataSize = atb::Utils::GetTensorSize(outTensors.at(0)); aclrtMalloc(&outTensors.at(i).deviceData, outTensors.at(i).dataSize, ACL_MEM_MALLOC_HUGE_FIRST); // Place the description and data size of other output tensors.
- Call the Setup API to calculate the workspace size.
uint64_t workspaceSize = 0; st = op->Setup(variantPack, workspaceSize, context);
- Allocate the NPU memory based on the workspace size.
void *workspace = nullptr; status = aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
If the workspace size is 0, skip this step. Otherwise, an error is reported.
- Call the Execute API to execute the operator.
st = op->Execute(variantPack, (uint8_t *)workspace, workspaceSize, context);
- Destroy the created object to release the memory.
status = aclrtDestroyStream(stream); // Destroy the stream. status = aclrtFree(workspace); // Destroy the workspace. st = atb::DestroyOperation(op); // Destroy the op object. st = atb::DestroyContext(context); // Destroy the context. // The following code is an example of releasing a tensor. In practice, all tensors in VariantPack need to be released. status = aclrtFree(tensor.deviceData); tensor.deviceData = nullptr; tensor.dataSize = 0;
# Use g++ to compile the demo project. demo.cpp is the source code file of the demo.
g++ -I "${ATB_HOME_PATH}/include" -I "${ASCEND_HOME_PATH}/include" -L "${ATB_HOME_PATH}/lib" -L "${ASCEND_HOME_PATH}/lib64" demo.cpp -l atb -l ascendcl -o demo
./demo # Run the executable file.
If abi=0, the compilation option -D_GLIBCXX_USE_CXX11_ABI=0 needs to be added to the g++ command.
Set the abi parameter for the ATB using either of the following methods if required:
- Automatic configuration: When the set_env.sh script is executed, if no parameter is added and the PyTorch environment is detected, the torch.compiled_with_cxx11_abi() API is automatically called, and the abi parameter during PyTorch compilation is selected as the abi parameter of the ATB. If the PyTorch environment is not detected, abi=1 is configured by default.
- Manual configuration: When the set_env.sh script is executed, you can specify the abi parameter of the ATB by using --cxx_abi=1 and --cxx_abi=0. Example: source set_env.sh --cxx_abi=1