CANN/ops-blas Iamax算子实现

📅 2026/7/15 1:55:06 👁️ 阅读次数 📝 编程学习
CANN/ops-blas Iamax算子实现

Iamax算子实现

【免费下载链接】ops-blas本项目是CANN提供的高性能线性代数计算以及轻量化GEMM调用算子库。项目地址: https://gitcode.com/cann/ops-blas

概述

BLAS Iamax算子实现。

Iamax(最大绝对值元素索引)算子实现了查找向量中绝对值最大的元素索引,是BLAS基础线性代数库中的核心算子之一。

该算子返回1-based索引,遵循BLAS惯例,常用于主元选择和迭代算法中。

支持的产品

  • Atlas A3 训练系列产品/Atlas A3 推理系列产品
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品

目录结构介绍

├── iamax │ ├── CMakeLists.txt // 编译工程文件 │ ├── README.md // 说明文档 │ └── iamax_test.cpp // 算子调用样例

算子描述

  • 算子功能:
    Iamax算子实现了查找向量中绝对值最大的元素索引。对应的数学表达式为:
result = argmax_i |x[i]|

返回1-based索引

对应的接口为:

int aclblasIamax(const float *x, int32_t *result, const int64_t n, const int64_t incx, const uint32_t dtypeFlag, void *stream);
参数iamax 参数说明
参数列表Param.Memoryin/out含义
nin向量元素个数。
xdevicein向量,包含 n 个元素。
incxinx 中连续元素之间的步长。
dtypeFlagin数据类型标志,0表示实数float,1表示复数float。
resultdeviceout最大绝对值元素的索引(1-based)。
  • 算子规格:

    算子类型(OpType)Iamax
    算子输入nameshapedata typeformat
    xnfloatND
    算子输出result1int32ND
    核函数名iamax_kernel
  • 算子实现:

    将输入数据从x的GM地址分块搬运到UB,并行计算各核的局部最大值,最后归约得到全局最大值索引。

  • 调用实现
    使用内核调用符<<<>>>调用核函数。

编译运行

在本样例根目录下执行如下步骤,编译并执行算子。

  • 配置环境变量
    请根据当前环境上CANN开发套件包的安装方式,选择对应配置环境变量的命令。

    • 默认路径,root用户安装CANN软件包

      source /usr/local/Ascend/cann/set_env.sh
    • 默认路径,非root用户安装CANN软件包

      source $HOME/Ascend/cann/set_env.sh
    • 指定路径install_path,安装CANN软件包

      source ${install_path}/cann/set_env.sh
  • 样例执行

    bash build.sh --ops=iamax --run # --ops=<算子名> --run可选参数,执行测试样例

    执行结果如下,说明精度对比成功。

    [Success] Case accuracy is verification passed.

【免费下载链接】ops-blas本项目是CANN提供的高性能线性代数计算以及轻量化GEMM调用算子库。项目地址: https://gitcode.com/cann/ops-blas

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考