0
代码
代码
Issues
Pull Requests
讨论
Wiki
分析
0
cann-robotcann-robotqli资料整改 & 增加examples
8e495e58创建于 23 天前历史提交

aclnnQuantLightningIndexer

📄 查看源码

产品支持情况

产品 是否支持
Ascend 950PR/Ascend 950DT √
Atlas A3 训练系列产品/Atlas A3 推理系列产品 ×
Atlas A2 训练系列产品/Atlas A2 推理系列产品 ×
Atlas 200I/500 A2 推理产品 ×
Atlas 推理系列产品 ×
Atlas 训练系列产品 ×

功能说明

  • 接口功能：QuantLightningIndexer在LightningIndexer的基础上支持了Per-Token-Head量化输入。

  • 计算公式：

    out=Top-k{[1]1×g@[(W@[1]1×Sk)⊙ReLU((ScaleQ@ScaleKT)⊙(QindexQuant@(KindexQuant)T))]}out = \text{Top-}k\left\{[1]_{1\times g}@\left[(W@[1]_{1\times S_{k}})\odot\text{ReLU}\left(\left(Scale_Q@Scale_K^T\right)\odot\left(Q_{index}^{Quant}@{\left(K_{index}^{Quant}\right)}^T\right)\right)\right]\right\}

函数原型

每个算子分为两段式接口，必须先调用“aclnnQuantLightningIndexerGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnQuantLightningIndexer”接口执行计算。

aclnnStatus aclnnQuantLightningIndexerGetWorkspaceSize(
    const aclTensor *query,
    const aclTensor *key,
    const aclTensor *weights,
    const aclTensor *queryDequantScale,
    const aclTensor *keyDequantScale,
    const aclTensor *actualSeqLengthsQueryOptional,
    const aclTensor *actualSeqLengthsKeyOptional,
    const aclTensor *blockTableOptional,
    int64_t          queryQuantMode,
    int64_t          keyQuantMode,
    char            *layoutQueryOptional,
    char            *layoutKeyOptional,
    int64_t          sparseCount,
    int64_t          sparseMode,
    int64_t          preTokens,
    int64_t          nextTokens,
    const aclTensor *out,
    uint64_t        *workspaceSize,
    aclOpExecutor   **executor)

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

aclnnQuantLightningIndexerGetWorkspaceSize

  • 参数说明：

Note

  • query、key、weights、query_dequant_scale、key_dequant_scale参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、N（Head Num）表示多头数、D（Head Dim）表示hidden层最小的单元尺寸，且满足D=H/N、T表示所有Batch输入样本序列长度的累加和。
  • 使用S1和S2分别表示query和key的输入样本序列长度，N1和N2分别表示query和key对应的多头数，k表示最后选取的索引个数。参数query中的D和参数key中的D值相等为128。T1和T2分别表示query和key的输入样本序列长度的累加和。
参数名 输入/输出 描述 使用说明 数据类型 数据格式 维度(shape) 非连续Tensor
query 输入 公式中的输入Q。 不支持空tensor。 INT8、FLOAT8_E4M3、HIFLOAT8。 ND
  • layout_query为BSND时，shape为(B,S1,N1,D)。
  • layout_query为TND时，shape为(T1,N1,D)。
 x
key 输入 公式中的输入K。
  • 不支持空tensor。
  • block_num为PageAttention时block总数，block_size为一个block的token数。
 INT8、FLOAT8_E4M3、HIFLOAT8。 ND
  • layout_key为PA_BSND时，shape为(block_num, block_size, N2, D)。
  • layout_key为BSND时，shape为(B, S2, N2, D)。
  • layout_key为TND时，shape为(T2, N2, D)。
 ✓
weights 输入 公式中的输入W。 不支持空tensor。 FLOAT16、BFLOAT16 ND
  • layout_query为BSND时，shape为(B,S1,N1)。
  • layout_query为TND时，shape为(T1,N1)。
 x
queryDequantScale 输入 表示Index Query的反量化系数Scale_Q。 不支持空tensor。 FLOAT、FLOAT16 ND
  • layout_query为BSND时，shape为(B,S1,N1)。
  • layout_query为TND时，shape为(T1,N1)。
 x
keyDequantScale 输入 表示Index Key的反量化系数Scale_K。 不支持空tensor。 FLOAT、FLOAT16 ND
  • layout_key为BSND时，shape为(B,S2,N2)。
  • layout_key为TND时，shape为(T2,N2)。
  • layout_key为PA_BSND时，shape为(block_num, block_size, N2)。
  • block_num为PageAttention时block总数，block_size为一个block的token数。
 ✓
actualSeqLengthsQueryOptional 输入 每个Batch中，Query的有效token数。
  • 不支持空tensor。
  • 如果不指定seqlen可传入None，表示和query的shape的S长度相同。
  • 该入参中每个Batch的有效token数不超过query中的维度S大小且不小于0，支持长度为B的一维tensor。
  • 当layout_query为TND时，该入参必须传入，且以该入参元素的数量作为B值，该入参中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。
  • 不能出现负值。
 INT32 ND (B,) x
actualSeqLengthsKeyOptional 输入 每个Batch中，Key的有效token数。
  • 不支持空tensor。
  • 如果不指定seqlen可传入None，表示和key的shape的S长度相同。
  • 该参数中每个Batch的有效token数不超过key/value中的维度S大小且不小于0，支持长度为B的一维tensor。
  • 当layout_key为TND或PA_BSND时，该入参必须传入，layout_key为TND，该参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。
  • 不能出现负值。
 INT32 ND (B,) x
blockTableOptional 输入 表示PageAttention中KV存储使用的block映射表。
  • 不支持空tensor。
  • PageAttention场景下，block_table必须为二维，第一维长度需要等于B，第二维长度不能小于maxBlockNumPerSeq（maxBlockNumPerSeq为每个batch中最大actual_seq_lengths_key对应的block数量）
  • block_size取值为16的整数倍，最大支持到1024。
 INT32 ND shape支持(B,S2_max/block_size) x
queryQuantMode 输入 用于标识输入query的量化模式。
  • 当前支持Per-Token-Head量化模式。
  • 当前仅支持传入0。
 INT64 - - -
keyQuantMode 输入 用于标识输入key的量化模式。
  • 当前支持Per-Token-Head量化模式。
  • 当前仅支持传入0。
 INT64 - - -
layoutQueryOptional 输入 用于标识输入Query的数据排布格式。 当前支持BSND、TND。 STRING - - -
layoutKeyOptional 输入 用于标识输入Key的数据排布格式。 当前支持PA_BSND、BSND、TND。 STRING - - -
sparseCount 输入 topK阶段需要保留的block数量。 当前支持[1, 2048]。 INT64 - - -
sparseMode 输入 表示sparse的模式。
  • sparse_mode为0时，代表defaultMask模式。
  • sparse_mode为3时，代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。
 INT64 - - -
preTokens 输入 用于稀疏计算，表示attention需要和前几个Token计算关联。 建议值2^63-1。 INT64 - - -
nextTokens 输入 用于稀疏计算，表示attention需要和后几个Token计算关联。 建议值2^63-1。 INT64 - - -
out 输出 公式中的Indices输出。 不支持空tensor。 INT32 ND
  • layout_query为"BSND"时输出shape为[B, S1, N2, sparseCount]。
  • layout_query为"TND"时输出shape为[T1, N2, sparseCount]。
 x
workspaceSize 输出 返回需要在Device侧申请的workspace大小。 - - - - -
executor 输出 返回op执行器，包含了算子计算流程。 - - - - -

  • 返回值：

    aclnnStatus：返回状态码，具体参见aclnn返回码。

    第一段接口会完成入参校验，出现以下场景时报错：

    返回值 错误码 描述
    ACLNN_ERR_PARAM_NULLPTR 161001 如果传入参数是必选输入，输出或者必选属性，且是空指针，则返回161001。
    ACLNN_ERR_PARAM_INVALID 161002 query、key、weights、queryDequantScale、keyDequantScale、actualSeqLengthsQueryOptional、actualSeqLengthsKeyOptional、queryQuantMode、keyQuantMode、layoutQueryOptional、layoutKeyOptional、sparseCount、sparseMode、out的数据类型和数据格式不在支持的范围内。

aclnnQuantLightningIndexer

  • 参数说明：

    参数名 输入/输出 描述
    workspace 输入 在Device侧申请的workspace内存地址。
    workspaceSize 输入 在Device侧申请的workspace大小，由第一段接口aclnnQuantLightningIndexerGetWorkspaceSize获取。
    executor 输入 op执行器，包含了算子计算流程。
    stream 输入 指定执行任务的Stream。

  • 返回值：

    aclnnStatus：返回状态码，具体参见aclnn返回码。

约束说明

  • 确定性说明：aclnnQuantLightningIndexer默认确定性实现。
  • 参数query中的N支持小于等于64/32/24/16，key的N支持1。
  • headdim支持128。
  • block_size取值为16的倍数，最大支持1024。
  • 参数query、key的数据类型应保持一致。
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品：
    • query和key的数据类型支持INT8。
    • 仅支持weights、query_dequant_scale、key_dequant_scale数据类型为FLOAT16、FLOAT16、FLOAT16。
  • Ascend 950PR/Ascend 950DT：
    • query N1仅支持8、16、24、32、64。
    • query和key的数据类型支持FLOAT8_E4M3、HIFLOAT8、INT8。
    • 当query和key的数据类型为FLOAT8_E4M3时，支持weights、query_dequant_scale、key_dequant_scale的数据类型为BFLOAT16、FLOAT、FLOAT或FLOAT16、FLOAT16、FLOAT16；
    • 当query和key的数据类型为HIFLOAT8时，仅支持weights、query_dequant_scale、key_dequant_scale数据类型为BFLOAT16、FLOAT、FLOAT；
    • 当query和key的数据类型为INT8时，仅支持weights、query_dequant_scale、key_dequant_scale数据类型为FLOAT16、FLOAT16、FLOAT16。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考编译与运行样例。

/**
 * Copyright (c) 2026 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file test_aclnn_quant_lightning_indexer.cpp
 * \brief
 */
//testci
#include <iostream>
#include <vector>
#include <cmath>
#include <cstring>
#include "securec.h"
#include "acl/acl.h"
#include "aclnnop/aclnn_quant_lightning_indexer.h"

using namespace std;

namespace {

#define CHECK_RET(cond) ((cond) ? true :(false))

#define LOG_PRINT(message, ...)     \
  do {                              \
    (void)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) {
  auto ret = aclInit(nullptr);
  if (!CHECK_RET(ret == ACL_SUCCESS)) {
    LOG_PRINT("aclInit failed. ERROR: %d\n", ret);
    return ret;
  }
  ret = aclrtSetDevice(deviceId);
  if (!CHECK_RET(ret == ACL_SUCCESS)) {
    LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret);
    return ret;
  }
  ret = aclrtCreateStream(stream);
  if (!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);
  auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
  if (!CHECK_RET(ret == ACL_SUCCESS)) {
    LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret);
    return ret;
  }

  ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
  if (!CHECK_RET(ret == ACL_SUCCESS)) {
    LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret);
    return ret;
  }

  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];
  }

  *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
  return 0;
}

struct TensorResources {
    void* queryDeviceAddr = nullptr;
    void* keyDeviceAddr = nullptr;
    void* weightsDeviceAddr = nullptr;
    void* queryDequantScaleDeviceAddr = nullptr;
    void* keyDequantScaleDeviceAddr = nullptr;
    void* actualSeqLengthsKeyDeviceAddr = nullptr;
    void* blockTableDeviceAddr = nullptr;
    void* outDeviceAddr = nullptr;

    aclTensor* queryTensor = nullptr;
    aclTensor* keyTensor = nullptr;
    aclTensor* weightsTensor = nullptr;
    aclTensor* queryDequantScaleTensor = nullptr;
    aclTensor* keyDequantScaleTensor = nullptr;
    aclTensor* actualSeqLengthsKeyTensor = nullptr;
    aclTensor* blockTableTensor = nullptr;
    aclTensor* outTensor = nullptr;
};

int InitializeTensors(TensorResources& resources) {
    std::vector<int64_t> queryShape = {2, 64, 16, 128};
    std::vector<int64_t> keyShape = {32, 16, 1, 128};
    std::vector<int64_t> weightsShape = {2, 64, 16};
    std::vector<int64_t> queryDequantScaleShape = {2, 64, 16};
    std::vector<int64_t> keyDequantScaleShape = {32, 16, 1};
    std::vector<int64_t> actualSeqLengthsKeyShape = {2};
    std::vector<int64_t> blockTableShape = {2, 32};
    std::vector<int64_t> outShape = {2, 64, 1, 2048};

    int64_t queryShapeSize = GetShapeSize(queryShape);
    int64_t keyShapeSize = GetShapeSize(keyShape);
    int64_t weightsShapeSize = GetShapeSize(weightsShape);
    int64_t queryDequantScaleShapeSize = GetShapeSize(queryDequantScaleShape);
    int64_t keyDequantScaleShapeSize = GetShapeSize(keyDequantScaleShape);
    int64_t actualSeqLengthsKeyShapeSize = GetShapeSize(actualSeqLengthsKeyShape);
    int64_t blockTableShapeSize = GetShapeSize(blockTableShape);
    int64_t outShapeSize = GetShapeSize(outShape);

    std::vector<int8_t> queryHostData(queryShapeSize, 1);
    std::vector<int8_t> keyHostData(keyShapeSize, 1);
    std::vector<uint16_t> weightsHostData(weightsShapeSize, 0x3F00);
    std::vector<float> queryDequantScaleHostData(queryDequantScaleShapeSize, 1.0f);
    std::vector<float> keyDequantScaleHostData(keyDequantScaleShapeSize, 1.0f);
    std::vector<int32_t> actualSeqLengthsKeyHostData = {256, 512};
    std::vector<int32_t> blockTableHostData(blockTableShapeSize, 0);
    for (int32_t i = 0; i < 32; i++) {
        blockTableHostData[i] = i;
        blockTableHostData[32 + i] = i;
    }
    std::vector<int32_t> outHostData(outShapeSize, 0);

    int ret = CreateAclTensor(queryHostData, queryShape, &resources.queryDeviceAddr,
                              aclDataType::ACL_FLOAT8_E4M3FN, &resources.queryTensor);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
      return ret;
    }

    ret = CreateAclTensor(keyHostData, keyShape, &resources.keyDeviceAddr,
                          aclDataType::ACL_FLOAT8_E4M3FN, &resources.keyTensor);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
      return ret;
    }

    ret = CreateAclTensor(weightsHostData, weightsShape, &resources.weightsDeviceAddr,
                          aclDataType::ACL_BF16, &resources.weightsTensor);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
      return ret;
    }

    ret = CreateAclTensor(queryDequantScaleHostData, queryDequantScaleShape, &resources.queryDequantScaleDeviceAddr,
                          aclDataType::ACL_FLOAT, &resources.queryDequantScaleTensor);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
      return ret;
    }

    ret = CreateAclTensor(keyDequantScaleHostData, keyDequantScaleShape, &resources.keyDequantScaleDeviceAddr,
                          aclDataType::ACL_FLOAT, &resources.keyDequantScaleTensor);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
      return ret;
    }

    ret = CreateAclTensor(actualSeqLengthsKeyHostData, actualSeqLengthsKeyShape, &resources.actualSeqLengthsKeyDeviceAddr,
                          aclDataType::ACL_INT32, &resources.actualSeqLengthsKeyTensor);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
      return ret;
    }

    ret = CreateAclTensor(blockTableHostData, blockTableShape, &resources.blockTableDeviceAddr,
                          aclDataType::ACL_INT32, &resources.blockTableTensor);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
      return ret;
    }

    ret = CreateAclTensor(outHostData, outShape, &resources.outDeviceAddr,
                          aclDataType::ACL_INT32, &resources.outTensor);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
      return ret;
    }
    return ACL_SUCCESS;
}

int ExecuteQuantLightningIndexer(TensorResources& resources, aclrtStream stream,
                                 void** workspaceAddr, uint64_t* workspaceSize) {
    int64_t queryQuantMode = 0;
    int64_t keyQuantMode = 0;
    int64_t sparseCount = 2048;
    int64_t sparseMode = 3;
    int64_t preTokens = 9223372036854775807;
    int64_t nextTokens = 9223372036854775807;
    constexpr const char layoutQueryStr[] = "BSND";
    constexpr const char layoutKeyStr[] = "PA_BSND";
    constexpr size_t layoutQueryLen = sizeof(layoutQueryStr);
    constexpr size_t layoutKeyLen = sizeof(layoutKeyStr);
    char layoutQuery[layoutQueryLen];
    char layoutKey[layoutKeyLen];
    errno_t memcpyRet = memcpy_s(layoutQuery, sizeof(layoutQuery), layoutQueryStr, layoutQueryLen);
    if (!CHECK_RET(memcpyRet == 0)) {
        LOG_PRINT("memcpy_s layoutQuery failed. ERROR: %d\n", memcpyRet);
        return -1;
    }
    memcpyRet = memcpy_s(layoutKey, sizeof(layoutKey), layoutKeyStr, layoutKeyLen);
    if (!CHECK_RET(memcpyRet == 0)) {
        LOG_PRINT("memcpy_s layoutKey failed. ERROR: %d\n", memcpyRet);
        return -1;
    }
    aclOpExecutor* executor;

    int ret = aclnnQuantLightningIndexerGetWorkspaceSize(
        resources.queryTensor, resources.keyTensor, resources.weightsTensor,
        resources.queryDequantScaleTensor, resources.keyDequantScaleTensor,
        nullptr, resources.actualSeqLengthsKeyTensor, resources.blockTableTensor,
        queryQuantMode, keyQuantMode, layoutQuery, layoutKey,
        sparseCount, sparseMode, preTokens, nextTokens,
        resources.outTensor, workspaceSize, &executor);

    if (!CHECK_RET(ret == ACL_SUCCESS)) {
        LOG_PRINT("aclnnQuantLightningIndexerGetWorkspaceSize failed. ERROR: %d\n", ret);
        return ret;
    }

    if (*workspaceSize > 0ULL) {
        ret = aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        if (!CHECK_RET(ret == ACL_SUCCESS)) {
            LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret);
            return ret;
        }
    }

    ret = aclnnQuantLightningIndexer(*workspaceAddr, *workspaceSize, executor, stream);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
        LOG_PRINT("aclnnQuantLightningIndexer failed. ERROR: %d\n", ret);
        return ret;
    }

    return ACL_SUCCESS;
}

int PrintOutResult(std::vector<int64_t> &shape, void** deviceAddr) {
  auto size = GetShapeSize(shape);
  std::vector<int32_t> resultData(size, 0);
  auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),
                         *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
  if (!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("sparse_indices result[%ld] is: %d\n", i, resultData[i]);
  }
  return ACL_SUCCESS;
}

void CleanupResources(TensorResources& resources, void* workspaceAddr,
                     aclrtStream stream, int32_t deviceId) {
    if (resources.queryTensor) {
      aclDestroyTensor(resources.queryTensor);
    }
    if (resources.keyTensor) {
      aclDestroyTensor(resources.keyTensor);
    }
    if (resources.weightsTensor) {
      aclDestroyTensor(resources.weightsTensor);
    }
    if (resources.queryDequantScaleTensor) {
      aclDestroyTensor(resources.queryDequantScaleTensor);
    }
    if (resources.keyDequantScaleTensor) {
      aclDestroyTensor(resources.keyDequantScaleTensor);
    }
    if (resources.actualSeqLengthsKeyTensor) {
      aclDestroyTensor(resources.actualSeqLengthsKeyTensor);
    }
    if (resources.blockTableTensor) {
      aclDestroyTensor(resources.blockTableTensor);
    }
    if (resources.outTensor) {
      aclDestroyTensor(resources.outTensor);
    }

    if (resources.queryDeviceAddr) {
      aclrtFree(resources.queryDeviceAddr);
    }
    if (resources.keyDeviceAddr) {
      aclrtFree(resources.keyDeviceAddr);
    }
    if (resources.weightsDeviceAddr) {
      aclrtFree(resources.weightsDeviceAddr);
    }
    if (resources.queryDequantScaleDeviceAddr) {
      aclrtFree(resources.queryDequantScaleDeviceAddr);
    }
    if (resources.keyDequantScaleDeviceAddr) {
      aclrtFree(resources.keyDequantScaleDeviceAddr);
    }
    if (resources.actualSeqLengthsKeyDeviceAddr) {
      aclrtFree(resources.actualSeqLengthsKeyDeviceAddr);
    }
    if (resources.blockTableDeviceAddr) {
      aclrtFree(resources.blockTableDeviceAddr);
    }
    if (resources.outDeviceAddr) {
      aclrtFree(resources.outDeviceAddr);
    }

    if (workspaceAddr) {
      aclrtFree(workspaceAddr);
    }
    if (stream) {
      aclrtDestroyStream(stream);
    }
    aclrtResetDevice(deviceId);
    aclFinalize();
}

} // namespace

int main() {
    int32_t deviceId = 0;
    aclrtStream stream = nullptr;
    TensorResources resources = {};
    void* workspaceAddr = nullptr;
    uint64_t workspaceSize = 0;
    std::vector<int64_t> outShape = {2, 64, 1, 2048};
    int ret = ACL_SUCCESS;

    // 1. Initialize device and stream
    ret = Init(deviceId, &stream);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
        LOG_PRINT("Init acl failed. ERROR: %d\n", ret);
        return ret;
    }

    // 2. Initialize tensors
    ret = InitializeTensors(resources);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
        CleanupResources(resources, workspaceAddr, stream, deviceId);
        return ret;
    }

    // 3. Execute the operation
    ret = ExecuteQuantLightningIndexer(resources, stream, &workspaceAddr, &workspaceSize);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
        CleanupResources(resources, workspaceAddr, stream, deviceId);
        return ret;
    }

    // 4. Synchronize stream
    ret = aclrtSynchronizeStream(stream);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
        LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret);
        CleanupResources(resources, workspaceAddr, stream, deviceId);
        return ret;
    }

    // 5. Process results
    PrintOutResult(outShape, &resources.outDeviceAddr);

    // 6. Cleanup resources
    CleanupResources(resources, workspaceAddr, stream, deviceId);
    return 0;
}