aclnnTopKTopPSampleV2

📄 View Source Code

Applicable Products

Function

• Description:

Performs topK-topP-minP-sample sampling computation based on the input logits, topK/topP/minP sampling parameters, and the random sampling weight distribution q. When isNeedSampleResult is set to false, outputs the maximum logits index logitsSelectIdx for each batch, along with the logits distribution logitsTopKPSelect after topK-topP-minP sampling. When isNeedSampleResult is set to true, outputs the intermediate computation results logitsIdx and logitsSortMasked after topK-topP-minP sampling, where logitsSortMasked is the intermediate result of logits after topK-topP-minP sampling computation, and logitsIdx is the corresponding index of logitsSortMasked in logits.

The operator includes four sampling algorithms that can be individually enabled but maintain a fixed upstream-to-downstream processing relationship (from original input to final output): TopK sampling, TopP sampling, MinP sampling, and exponential sampling (referred to as Sample in this document). The following computation scenarios are currently supported, as shown in the table below:

• Formula:

The input logits is a word frequency table of size [batch, voc_size], where each batch corresponds to an input sequence, and voc_size is the uniform length agreed upon for each batch.
Each row logits[batch][:] in logits executes different computation scenarios based on the corresponding topK[batch], topP[batch], minP[batch, :], and q[batch, :].
In the formulas below, b and v are used to denote the indices along the batch and voc_size dimensions, respectively.

TopK Sampling

1. Uses segmented TopK merge sort with segment length v, where the TopK of block {s-1} is used to pre-filter the input of the current block {s}, progressively updating the TopK of a single batch to reduce redundant data and computation.

2. topK[batch] corresponds to the k value for sampling the current batch, with a valid range of 1 ≤ topK[batch] ≤ min(voc_size[batch], ks_max). If top[k] exceeds the valid range, the TopK sampling stage for the current batch is skipped, and the sorting for the current batch is also skipped, with the input logits[batch] passed directly to the next module.

• Split the current batch into several sub-segments and compute topKValue[b] in a rolling manner:

$$topKValue[b] = {Max(topK[b])}_{s=1}^{\left \lceil \frac{S}{v} \right \rceil }\left \{ topKValue[b]\left \{s-1 \right \} \cup \left \{ logits[b][v] \ge topKMin[b][s-1] \right \} \right \}\\ Card(topKValue[b])=topK[b]$$

Where:

$$topKMin[b][s] = Min(topKValue[b]\left \{ s \right \})$$

v denotes the preset fixed segment length for rolling topK:

$$v = 8 * \text{ks\_max}$$

The valid range of ks_max is [1, 1024], defaults to 1024, and must be aligned upward to a multiple of 8.

• Generate the mask to be filtered.

$$sortedValue[b] = sort(topKValue[b], descendant)$$ $$topKMask = sortedValue \geq topKValue$$

• Set the parts below the threshold to defLogit via the mask:

$$sortedValue[b][v]= \begin{cases} defLogit & \text{topKMask[b][v] = false} \\ sortedValue[b][v] & \text{topKMask[b][v] = true} & \end{cases}$$

• defLogit depends on the input constraint attribute input_is_logits, which controls the normalization of the input logits and the output logits_top_kp_select:

$$\text{defLogit} = \begin{cases} -inf, & \text{inputIsLogits} = \text{true} \\ 0, & \text{inputIsLogits} = \text{false} \end{cases}$$

TopP Sampling

• Based on the input constraint attribute inputIsLogits, if this attribute is True, the sorted results are normalized:

  $$\text{logit\_sortProb} = \begin{cases} \text{softmax}(\text{logits\_sort}), & \text{inputIsLogits} = \text{True} \\ \text{logits\_sort}, & \text{inputIsLogits} = \text{False} \end{cases}$$

• Based on the value of the input top_p[b], this module adopts the following processing strategies:

  Parameter Type	≤0	Valid Range	Invalid Range
  top_p[b]	Retains one token with the highest logit.	0 < top_p < 1, performs topP sampling.	top_p ≥ 1, skips topP sampling.

• If regular topP sampling is performed and the preceding topK stage has already produced sorted output results, the cumulative logit probabilities are computed based on the topK sampling output, and sampling is truncated according to top_p:

  $$topPMask[b] = \begin{cases} 0, & \sum_{\text{topKMask}[b]}^{} \text{logits\_sortProb}[b][*] > p[b] \\ 1, & \sum_{\text{topKMask}[b]}^{} \text{logits\_sortProb}[b][*] \leq p[b] \end{cases}$$

• If regular topP sampling is performed but the preceding topK stage is skipped, the top-p mask is computed:

  $$topPMask[b] = \begin{cases} topKMask[b][0:GuessK], & \sum_{\text{GuessK}}^{} probValue[b][*] \ge p[b] \\ probSum[b][v] \le 1 - p[b], & \text{others} \end{cases}$$

• Set the positions to be filtered to the default invalid value defLogit to obtain logits_sort, denoted as sortedValue[b][v]:

$$sortedValue[b][v] = \begin{cases} defLogit & \quad \text{topPMask}[b][v] = \text{false} \\ logit\_sortProb[b][v] & \quad \text{topPMask}[b][v] = \text{true} \end{cases}$$

• Take the first topK elements in each row of the filtered sortedValue[b][v], look up their original indices in the input, and combine them into logits_idx:

$$logitsIdx[b][v] = Index(sortedValue[b][v] \in Logits)$$

• Use the truncated sortedValue as logitsSortMasked:

$$logitsSortMasked[b,:] = sortedValue[b]$$

min-p sampling

• If min_ps[b] ∈ (0, 1), perform min-p sampling:

  $$\text{logitsMax}[b] = \text{Max}(\text{logitsSortMasked}[b])$$ $$\text{minPThd} = \text{logitsMax}[b] * \text{minPs}[b]$$ $$\text{minPMask}[b] = \begin{cases} 0, & \text{logitsSortMasked}[b] < \text{minPThd} \\ 1, & \text{logitsSortMasked}[b] \geq \text{minPThd} \end{cases}$$ $$\text{logitsSortMasked}[b,:] = \begin{cases} \text{defLogit}, & \text{minPMask}[b] = 0 \\ \text{logitsSortMasked}[b,:], & \text{minPMask}[b] = 1 \end{cases}$$

• Other cases:

  $$\text{logitsSortMasked}[b, :] = \begin{cases} \text{logitsSortMasked}[b, :], & \text{if } minPs[b] \leq 0 \\ \max(\text{logitsSortMasked}[b, :]), & \text{if } minPs[b] \geq 1 \end{cases}$$

  When min_ps[b] ≥ 1, only one maximum token is taken per batch, and the remaining positions are filled with defLogit.

Optional Outputs

• If the input attribute IsNeedLogits is set to True, logitsIndexMasked after combined topK-topP-minP sampling is used for logits_top_kp_select output.

  $$\text{logitsIndex}[b][v] = \text{Index}(\text{logitsSortMasked}[b][v] \in \text{Logits})$$ $$\text{logitsIndexMasked}[b,:] = \text{logitsIndex}[b,:] * \text{topKMask}[b] * \text{topPMask}[b] * \text{minPMask}[b]$$

  If the topK, topP, or minP sampling step is skipped, the corresponding mask is all ones.

• Next, logitsIndexMasked is used to select from the input Logits, filtering high-frequency tokens in the input Logits as the logits_top_kp_select output:

  $$\text{logitsTopKpSelect}[b][v] = \begin{cases} \text{logits}[b][v], & \text{if } logitsIndexMasked[b,v] = \text{True} \\ \text{defLogit}, & \text{if } logitsIndexMasked[b,v] = \text{False} \end{cases}$$

Subsequent Processing

• The input at this stage is logitsSortMasked, which is the combined result of the preceding topK-topP-minP sampling.

• The input here must satisfy logitsSortMasked ∈ (0,1). Configure the input constraint attribute inputIsLogits based on the actual input Logits, as follows:

  $$\text{inputIsLogits} = \begin{cases} True, & \text{Logits} \notin [0,1] \\ False, & \text{Logits} \in [0,1] \end{cases}$$

  such that

  $$\text{probs}[b] = \text{logitsSortMasked}[b, :]$$

  There are three subsequent modes: None, QSample, and outputting intermediate results, controlled by the input constraint attribute isNeedSampleResult and whether q is passed.

• None:

• This mode is used when isNeedSampleResult is false and q is not provided. In this mode, the maximum element and index are directly obtained for each batch through Argmax and output via gatherOut.

  $$\text{logitsSelectIdx}[b] = \text{LogitsIdx}[b]\left[\text{ArgMax}(\text{probs}[b][:])\right]$$

• QSample:

• This mode is used when isNeedSampleResult is false and q is provided. In this mode, exponential distribution sampling is first performed on probs:

  $$qCnt = \text{Sum}(\text{MinPMask} == 1)$$ $$\text{probsOpt}[b] = \frac{\text{probs}[b]}{q[b, :qCnt] + \text{eps}}$$

• Then perform Argmax-GatherOut to output the result:

  $$\text{logitsSelectIdx}[b] = \text{LogitsIdx}[b][\text{ArgMax}(\text{probsOpt}[b][:])]$$

• Output intermediate results:

• When isNeedSampleResult is true, this mode is used. In this case, the sampled logitsSortMasked and its original index logitsIdx in the input are output:

  $$\text{logitsSortMasked}[b, v] = \begin{cases} \text{logitsSortMasked}[b, v], & \text{if } \text{minPMask}[b, v] = 1 \\ 0, & \text{if } \text{minPMask}[b, v] = 0 \end{cases}$$ $$logitsIdx[b][v] = Index(logitsSortMasked[b][v])$$

Function Prototype

Each operator is divided into a two-phase API. The aclnnTopKTopPSampleV2GetWorkspaceSize API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the aclnnTopKTopPSampleV2 API is called to perform the computation.

aclnnStatus aclnnTopKTopPSampleV2GetWorkspaceSize(
  const aclTensor *logits,
  const aclTensor *topK,
  const aclTensor *topP,
  const aclTensor *q,
  const aclTensor *minPs,
  double           eps,
  bool             isNeedLogits,
  int64_t          topKGuess,
  int64_t          ksMax,
  bool             inputIsLogits,
  bool             isNeedSampleResult,
  const aclTensor *logitsSelectIdx,
  const aclTensor *logitsTopKPSelect,
  const aclTensor *logitsIdx,
  const aclTensor *logitsSortMasked,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

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

aclnnTopKTopPSampleV2GetWorkspaceSize

• Parameters:

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  logits	Input	Represents the input logits to be sampled, with the logits index fixed to the last dimension, corresponding to `logits` in the formula.	Empty tensors are not supported.	FLOAT16, BFLOAT16, FLOAT32	ND	2	√
  topK	Input	Indicates the k value for sampling each batch, corresponding to `topK[b]` in the formula.	Empty tensors are not supported. The shape must be consistent with the first n-1 dimensions of logits.	INT32	ND	1	√
  topP	Input	Indicates the p value for sampling each batch. Corresponds to `topP[b]` in the formula.	Does not support empty tensors. The shape must be consistent with the first n-1 dimensions of `logits`, and the data type must be consistent with `logits`.	FLOAT16, BFLOAT16, FLOAT32	ND	1	√
  q	Input	Represents the exponential sampling matrix output by topK-topP-minP sampling. Corresponds to `q` in the formula.	The shape must be consistent with `logits`.	FLOAT32	ND	2	√
  minPs	Input	Represents the minP value for sampling in each batch. Corresponds to `minPs[b]` in the formula.	The shape must be consistent with the first n-1 dimensions of `logits`, and the data type must be consistent with `logits`.	FLOAT16, BFLOAT16, FLOAT32	ND	2	√
  eps	Input	Prevents division by zero in softmax and weighted sampling. The recommended value is 1e-8.	-	DOUBLE	-	-	-
  isNeedLogits	Input	Controls the output condition of logitsTopKPSelect. It is recommended to set this to 0.	-	BOOL	-	-	-
  topKGuess	Input	Indicates the candidate logits size for each batch when attempting topP partial traversal sampling of logits.	-	INT64	-	-	-
  ksMax	Input	Specifies the maximum topK value for each batch during topK sampling. Must be a positive integer.	-	INT64	-	-	-
  inputIsLogits	Input	Indicates whether the input logits are unnormalized. Defaults to true.	-	BOOL	-	-	-
  isNeedSampleResult	Input	Specifies whether to output intermediate computation results. Defaults to false.	-	BOOL	-	-	-
  logitsSelectIdx	Output	Indicates the position index of the element with the maximum frequency max(probsOpt[batch, :]) in each batch within the input logits after the topK-topP-minP-sample computation process.	Empty tensors are not supported. The shape must be consistent with the first n-1 dimensions of logits.	INT64	ND	1	√
  logitsTopKPSelect	Output	Indicates the logits that remain unfiltered from the input logits after the topK-topP-minP computation process.	Empty tensors are not supported. The shape must be consistent with logits.	FLOAT32	ND	1	√
  logitsIdx	Output	Indicates the position index of the intermediate sampling result for each batch in the input logits after the topK-topP-minP computation pipeline.	Empty tensors are not supported. The shape must be the same as that of logits.	INT64	ND	1	√
  logitsSortMasked	Output	Indicates the intermediate sampling result of each batch after the topK-topP-minP computation process.	Does not support empty tensors. The shape must be the same as that of logits.	FLOAT32	ND	1	√
  workspaceSize	Output	Returns the workspace size to be requested on the device side.	-	-	-	-	-
  executor	Output	Returns the operator executor, which contains the operator computation flow.	-	-	-	-	-

• Return value

  aclnnStatus: return code. For details, see aclnn Return Code.

  The first-phase API performs input parameter validation and returns an error in the following scenarios:

  Return Code	Error Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	Any of the input parameters logits, topK, or topP is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data type of logits, topK, topP, q, or minPs is not within the supported range.
  		The dimensions or sizes of logits and q are inconsistent.
  		The dimensions of topK, topP, and minPs are inconsistent with the first n-1 dimensions of logits.
  		The data types of logits, topP, and minPs are inconsistent.

aclnnTopKTopPSampleV2

• Parameters:

  Parameter	Input/Output	Description
  workspace	Input	Specifies the workspace memory address applied for on the Device side.
  workspaceSize	Input	Specifies the workspace size applied for on the Device side, which is obtained by the first-phase API aclnnTopKTopPSampleV2GetWorkspaceSize.
  executor	Input	Operator executor that contains the operator computation flow.
  stream	Input	Specifies the stream for task execution.

• Return value

  aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

• Deterministic computation:

  • aclnnTopKTopPSampleV2 defaults to a deterministic implementation.

• For all sampling parameters, their dimensions must satisfy batch > 0 and 0 < vocSize <= 2^20.

• topK accepts only non-negative values as valid input. Passing 0 or a negative value skips sampling for the corresponding batch.

• The sizes and dimensions of logits, q, logitsTopKPselect, logitsIdx, and logitsSortMasked must be exactly the same.

• All dimensions except the last one of logits, topK, topP, minPs, logitsSelectIdx, logitsIdx, and logitsSortMasked must be exactly the same in order and size. Currently, logits can only be 2-dimensional, and topK, topP, and logitsSelectIdx must be 1-dimensional non-empty tensors. logits, topK, and topP do not allow empty tensors as input. To skip the corresponding module, set the input according to the corresponding rules.

• To skip the topK module individually, pass a tensor of size [batch, 1] with each element set to an invalid value.

• If min(ksMaxAligned, 1024) < topK[batch] < vocSize[batch], all valid elements of the current batch are selected and topK sampling is skipped. ksMaxAligned is ksMax rounded up to the nearest multiple of 8, and the value range of ksMax is [1, 1024].

• To skip the topP module individually, pass a tensor of size [batch, 1] with each element ≥ 1.

• To skip the minP module, pass minPs=nullptr or pass a tensor of size [batch, 1] with each element ≤ 0.

• To skip the sample module, pass q=nullptr. To use the sample module, pass a tensor of size [batch, vocSize].

• To output intermediate results, set isNeedSampleResult to true and pass q=nullptr. In this case, logitsSelectIdx is not output.

Example

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_top_k_top_p_sample_v2.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 for resource initialization.
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 device memory.
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);
// Copy host-side data to device memory by calling aclrtMemcpy.
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);

// Calculate the strides of a 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/stream. See the ACL API.
  // Set deviceId based on the actual device.
  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 input and output. Customize the construction based on the API.
  std::vector<int64_t> logitsShape = {48, 131072};
  std::vector<int64_t> topKPShape = {48};
  long long vocShapeSize = GetShapeSize(logitsShape);
  long long batchShapeSize = GetShapeSize(topKPShape);

  void* logitsDeviceAddr = nullptr;
  void* topKDeviceAddr = nullptr;
  void* topPDeviceAddr = nullptr;
  void* qDeviceAddr = nullptr;
  void* minPsDeviceAddr = nullptr;
  void* logitsSelectedIdxDeviceAddr = nullptr;
  void* logitsTopKPSelectDeviceAddr = nullptr;
  void* logitsIdxDeviceAddr = nullptr;
  void* logitsSortMaskedDeviceAddr = nullptr;

  aclTensor* logits = nullptr;
  aclTensor* topK = nullptr;
  aclTensor* topP = nullptr;
  aclTensor* q = nullptr;
  aclTensor* minPs = nullptr;
  aclTensor* logitsSelectedIdx = nullptr;
  aclTensor* logitsTopKPSelect = nullptr;
  aclTensor* logitsIdx = nullptr;
  aclTensor* logitsSortMasked = nullptr;
  std::vector<int16_t> logitsHostData(48 * 131072, 1);
  std::vector<int32_t> topKHostData(48, 128);
  std::vector<int16_t> topPHostData(48, 1);
  std::vector<float> qHostData(48 * 131072, 1.0f);
  std::vector<int16_t> minPsHostData(48, 1);

  std::vector<int64_t> logitsSelectedIdxHostData(48, 0);
  std::vector<float> logitsTopKPSelectHostData(48 * 131072, 0);
  std::vector<int64_t> logitsIdxHostData(48 * 131072, 0);
  std::vector<float> logitsSortMaskedHostData(48 * 131072, 0);

  float eps = 1e-8;
  int64_t isNeedLogits = 0;
  int32_t topKGuess =32;
  int32_t ks_max = 1024;
  bool inputIsLogits = true;
  bool isNeedSampleResult = false;

  // Create the logits aclTensor.
  ret = CreateAclTensor(logitsHostData, logitsShape, &logitsDeviceAddr, aclDataType::ACL_BF16, &logits);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the topK aclTensor.
  ret = CreateAclTensor(topKHostData, topKPShape, &topKDeviceAddr, aclDataType::ACL_INT32, &topK);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the topP aclTensor.
  ret = CreateAclTensor(topPHostData, topKPShape, &topPDeviceAddr, aclDataType::ACL_BF16, &topP);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the q aclTensor.
  ret = CreateAclTensor(qHostData, logitsShape, &qDeviceAddr, aclDataType::ACL_FLOAT, &q);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the minps aclTensor.
  ret = CreateAclTensor(minPsHostData, topKPShape, &minPsDeviceAddr, aclDataType::ACL_BF16, &minPs);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the logitsSelected aclTensor.
  ret = CreateAclTensor(logitsSelectedIdxHostData, topKPShape, &logitsSelectedIdxDeviceAddr, aclDataType::ACL_INT64, &logitsSelectedIdx);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the logitsTopKPSelect aclTensor.
  ret = CreateAclTensor(logitsTopKPSelectHostData, logitsShape, &logitsTopKPSelectDeviceAddr, aclDataType::ACL_FLOAT, &logitsTopKPSelect);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the logitsIdx aclTensor.
  ret = CreateAclTensor(logitsIdxHostData, logitsShape, &logitsIdxDeviceAddr, aclDataType::ACL_INT64, &logitsIdx);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the logitsSortMasked aclTensor.
  ret = CreateAclTensor(logitsSortMaskedHostData, logitsShape, &logitsSortMaskedDeviceAddr, aclDataType::ACL_FLOAT, &logitsSortMasked);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // 3. Call the CANN operator library API. Replace with the actual API name.
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;
  // Call the first-phase API of aclnnTopKTopPSampleV2.
  ret = aclnnTopKTopPSampleV2GetWorkspaceSize(logits, topK, topP, q, minPs, eps, isNeedLogits, topKGuess, ks_max, inputIsLogits,
    isNeedSampleResult, logitsSelectedIdx, logitsTopKPSelect, logitsIdx, logitsSortMasked, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTopKTopPSampleV2GetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // Apply for device memory based on the 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 aclnnTopKTopPSampleV2.
  ret = aclnnTopKTopPSampleV2(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTopKTopPSampleV2 failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Synchronize and wait for task execution to 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 side memory to the host side. Modify based on the specific API definition.
  auto size = GetShapeSize(topKPShape);
  std::vector<int64_t> resultData(size, 0);
  ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), logitsSelectedIdxDeviceAddr,
                      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[%ld] is: %ld\n", i, resultData[i]);
  }

  // 6. Destroy the aclTensor. Modify based on the specific API definition.
  aclDestroyTensor(logits);
  aclDestroyTensor(topK);
  aclDestroyTensor(topP);
  aclDestroyTensor(q);
  aclDestroyTensor(logitsSelectedIdx);
  aclDestroyTensor(logitsTopKPSelect);
  // 7. Release device resources. Modify based on the specific API definition.
  aclrtFree(logitsDeviceAddr);
  aclrtFree(topKDeviceAddr);
  aclrtFree(topPDeviceAddr);
  aclrtFree(qDeviceAddr);
  aclrtFree(logitsSelectedIdxDeviceAddr);
  aclrtFree(logitsTopKPSelectDeviceAddr);
  if (workspaceSize > 0) {
      aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  return 0;
  }