Single-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.

1. Include the ACL and ATB API header files.

   #include <acl/acl.h>
   #include <atb/atb_infer.h>

2. Set deviceId.
   Set deviceId as required.

   int deviceId = 0;
   aclError status = aclrtSetDevice(deviceId);

3. 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);

4. Create a single-operator (OpsOperation) object instance.
   1. Construct an Operation parameter.
      Instantiate the parameter structure based on the operator to be created. For details about the API definition of the parameter structure, see atb/infer_op_params.h and atb/train_op_params.h.

      // Take the Add operator in the elewise category as an example. You can use the following method to construct the corresponding parameter:
      atb::infer::ElewiseParam param;
      param.elewiseType = atb::infer::ElewiseParam::ELEWISE_ADD;

   2. Create an operator object instance.

      atb::Operation *op = nullptr;
      atb::Status st = atb::CreateOperation(param, &op);

5. 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;
   variantPack.inTensors = { a, ... };
   variantPack.outTensors = { output, ... };

6. Call the Setup API to calculate the workspace size.

   uint64_t workspaceSize = 0;
   st = op->Setup(variantPack, workspaceSize, context);

7. 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.

8. Call the Execute API to execute the operator.

   st = op->Execute(variantPack, (uint8_t *)workspace, workspaceSize, context);

9. 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;

To run the demo, perform the following steps:

# 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