aclnnSort

Supported Products

Product	Supported
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	√
Atlas training products	√

Function

Sorts elements in the input tensor in ascending or descending order along a specified dimension and returns the corresponding index value. The input tensor self is N-dimensional [0, N – 1] and is sorted based on the dimension specified by dim.

Prototype

Each operator has two-phase API calls. First, aclnnSortGetWorkspaceSize is called to obtain the workspace size required for computation and the executor that contains the operator computation process. Then, aclnnSort is called to perform computation.

aclnnStatus aclnnSortGetWorkspaceSize(
  const aclTensor *self, 
  bool             stable, 
  int64_t          dim, 
  bool             descending, 
  aclTensor       *valuesOut, 
  aclTensor       *indicesOut, 
  uint64_t        *workspaceSize, 
  aclOpExecutor  **executor)

aclnnStatus aclnnSort(
  void             *workspace, 
  uint64_t          workspaceSize, 
  aclOpExecutor    *executor, 
  const aclrtStream stream)

aclnnSortGetWorkspaceSize

Parameters:

Name	Input/Output	Description	Usage	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
self	Input	Input tensor.	-	BFLOAT16, FLOAT16, FLOAT32, INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64, or BOOL	ND	1-8	√
stable	Input	Whether the sorting is stable. True: stable; False: unstable.	-	BOOL	-	-	-
dim	Input	Dimension along which the sorting is performed.	The range is [–self.dim(), self.dim() – 1]. self.dim() is the dimension of the input tensor.	INT64	-	-	-
descending	Input	Controls the sorting order. `True` indicates descending order, and `False` indicates ascending order.	-	BOOL	-	-	-
valuesOut	Output	Result of sorting the tensor along the specified dimension.	-	BFLOAT16, FLOAT16, FLOAT32, DOUBLE, INT8, INT16, INT32, INT64, UINT8, BOOL, UINT16, UINT32, or UINT64	ND	Same as that of self.	√
indicesOut	Output	Index of each element in the original tensor after sorting.	-	INT64	ND	Same as that of self	√
workspaceSize	Output	Size of the workspace required to be allocated on the device.	-	-	-	-	-
executor	Output	Operator executor, containing the operator computation process.	-	-	-	-	-

Atlas inference products and Atlas training products:
- The data type of self cannot be BFLOAT16, UINT16, UINT32, or UINT64.
- The data type of valuesOut cannot be BFLOAT16, UINT16, UINT32, or UINT64.
Atlas A2 training products/Atlas A2 inference products and Atlas A3 training products/Atlas A3 inference products:
- The data type of self cannot be UINT16, UINT32, or UINT64.
- When the data type of self is BFLOAT16, the axis specified by the dim parameter cannot be 1.
- The data type of valuesOut cannot be UINT16, UINT32, or UINT64.

Returns:

aclnnStatus: status code. For details, see aclnn Return Codes.

The first-phase API implements input parameter verification. The following errors may be thrown:

Return	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	The input `self`, `valuesOut`, or `indicesOut` is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type of `self`, `valuesOut`, or `indicesOut` is not supported, or the shapes do not match.
ACLNN_ERR_PARAM_INVALID	161002	The value of `dim` is not within the dimension range of the input tensor `self`.

aclnnSort

Parameters:

Name	Input/Output	Description
workspace	Input	Address of the workspace to be allocated on the device.
workspaceSize	Input	Size of the workspace to be allocated on the device, which is obtained by calling the first-phase API `aclnnSortGetWorkspaceSize`.
executor	Input	Operator executor, containing the operator computation process.
stream	Input	Stream for executing the task.

Returns:

aclnnStatus: status code. For details, see aclnn Return Codes.

Constraints

Deterministic compute:
- aclnnSort defaults to a deterministic implementation.
Atlas training products: If the input is FLOAT32, it will be converted to FLOAT16 for sorting and then converted back to FLOAT32.
If the data type of self is not FLOAT, FLOAT16, or BFLOAT16, the operator execution may time out due to an overlarge tensor size (an AI CPU error is reported, with reason=[aicpu timeout]). The maximum size for each data type (closely related to the remaining memory of the machine) is as follows:
- INT64 type: 150000000
- UINT8, INT8, INT16, INT32 types: 725000000

Example

The following example is for reference only. For details, see Compilation and Running Sample.

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_sort.h"

#define CHECK_RET(cond, return_expr) \
  do {                               \
    if (!(cond)) {                   \
      return_expr;                   \
    }                                \
  } while (0)

#define LOG_PRINT(message, ...)     \
  do {                              \
    printf(message, ##__VA_ARGS__); \
  } while (0)

int64_t GetShapeSize(const std::vector<int64_t>& shape) {
  int64_t shapeSize = 1;
  for (auto i : shape) {
    shapeSize *= i;
  }
  return shapeSize;
}

int Init(int32_t deviceId, aclrtStream* stream) {
  // (Boilerplate) Initialize resources.
  auto ret = aclInit(nullptr);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
  ret = aclrtSetDevice(deviceId);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
  ret = aclrtCreateStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
  return 0;
}

template <typename T>
int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
                    aclDataType dataType, aclTensor** tensor) {
  auto size = GetShapeSize(shape) * sizeof(T);
  // Call aclrtMalloc to allocate memory on the device.
  auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
  // Call aclrtMemcpy to copy the data on the host to the memory on the device.
  ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);

  // Compute the strides of the contiguous tensor.
  std::vector<int64_t> strides(shape.size(), 1);
  for (int64_t i = shape.size() - 2; i >= 0; i--) {
    strides[i] = shape[i + 1] * strides[i + 1];
  }

  // Call aclCreateTensor to create an aclTensor.
  *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr);
  return 0;
}

int main() {
  // 1. (Boilerplate) Initialize the device and stream. For details, see the ACL API manual.
  // Set the device ID in use.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

  // 2. Construct the inputs and outputs based on the API definition.
  bool stable = false;
  int64_t dim = 0;
  bool descending = false;
  std::vector<int64_t> selfShape = {3, 4};
  std::vector<int64_t> outValuesShape = {3, 4};
  std::vector<int64_t> outIndicesShape = {3, 4};
  void* selfDeviceAddr = nullptr;
  void* outValuesDeviceAddr = nullptr;
  void* outIndicesDeviceAddr = nullptr;
  aclTensor* self = nullptr;
  aclTensor* outValues = nullptr;
  aclTensor* outIndices = nullptr;
  std::vector<int64_t> selfHostData = {7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6};
  std::vector<int64_t> outValuesHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
  std::vector<int64_t> outIndicesHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

  // Create a self aclTensor.
  ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_INT64, &self);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create outValues and outIndices aclTensor.
  ret = CreateAclTensor(outValuesHostData, outValuesShape, &outValuesDeviceAddr, aclDataType::ACL_INT64, &outValues);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(outIndicesHostData, outIndicesShape, &outIndicesDeviceAddr, aclDataType::ACL_INT64, &outIndices);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // 3. Call the CANN operator library API, which needs to be replaced with the actual API.
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;
  // Call the first-phase API of aclnnSort.
  ret = aclnnSortGetWorkspaceSize(self, stable, dim, descending, outValues, outIndices, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSortGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // Allocate device memory based on workspaceSize computed by the first-phase API.
  void* workspaceAddr = nullptr;
  if (workspaceSize > 0) {
    ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
  }
  // Call the second-phase API of aclnnSort.
  ret = aclnnSort(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSort failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Wait until the task execution is complete.
  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  // 5. Obtain the output value and copy the result from the device memory to the host. Modify the code based on the API definition.
  auto size = GetShapeSize(outValuesShape);
  std::vector<int64_t> resultData(size, 0);
  ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outValuesDeviceAddr,
                    size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
  for (int64_t i = 0; i < size; i++) {
    LOG_PRINT("result values [%ld] is: %ld\n", i, resultData[i]);
  }

  auto size2 = GetShapeSize(outIndicesShape);
  std::vector<int64_t> resultData2(size2, 0);
  ret = aclrtMemcpy(resultData2.data(), resultData2.size() * sizeof(resultData2[0]), outIndicesDeviceAddr,
                    size * sizeof(resultData2[0]), ACL_MEMCPY_DEVICE_TO_HOST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
  for (int64_t i = 0; i < size2; i++) {
    LOG_PRINT("result indices [%ld] is: %ld\n", i, resultData2[i]);
  }

  // 6. Release aclTensor and aclScalar. Modify the code based on the API definition.
  aclDestroyTensor(self);
  aclDestroyTensor(outValues);
  aclDestroyTensor(outIndices);

  // 7. Release device resources. Modify the code based on the API definition.
  aclrtFree(selfDeviceAddr);
  aclrtFree(outValuesDeviceAddr);
  aclrtFree(outIndicesDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  return 0;
}