尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

使用 pyasc 获取 GlobalTensor 元素个数:`get_size` 接口解析与实战指南

使用 pyasc 获取 GlobalTensor 元素个数:`get_size` 接口解析与实战指南 使用 pyasc 获取 GlobalTensor 元素个数get_size接口解析与实战指南【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc在 CANN pyasc 项目中asc.language.core.GlobalTensor是面向 Python 开发者、用于承载昇腾 AI 处理器 Global Memory全局存储数据的核心 Tensor 类型其全部成员函数与 Ascend C 的GlobalTensor接口一一对应。本篇指南以 GlobalTensor.get_size 文档为主线系统讲解该接口的语义、底层实现、初始化前提、约束限制及实际算子中的典型用法。阅读完本文后你将掌握如何在asc.jit算子内核中正确获取全局 Tensor 的元素个数并能区分get_size与get_shape_info、set_global_buffer之间的关系避免常见误用。接口总览GlobalTensor.get_size()的定位函数签名与语义get_size是GlobalTensor的成员方法用于获取 GlobalTensor 的元素个数element count而非字节数。其 Python 侧签名如下GlobalTensor.get_size() → int参数无。返回值GlobalTensor的元素个数类型为int在内核编译期对应 64 位无符号整数uint64_t。对应 Ascend C 函数原型__aicore__ inline uint64_t GetSize() const从 C 原型可以看出该接口在设备侧AI Core以内联函数形式实现返回uint64_t与 Python 侧int语义一致。在 pyasc 中GlobalTensor位于 python/asc/language/core/tensor.py与LocalTensorLocal Memory 数据形成对照GlobalTensor专门存放全局数据类型T支持基础数据类型以及TensorTrait类型但需遵循使用该GlobalTensor的指令如data_copy、向量运算等的数据类型支持情况。与LocalTensor.get_size的区别pyasc 中LocalTensor同样提供get_size方法见 LocalTensor.get_size 文档但二者有本质区别对比项GlobalTensor.get_sizeLocalTensor.get_size存储位置Global Memory外部存储Local MemoryAI Core 内部存储返回值语义元素个数当前 Size 大小单位为元素返回类型uint64_tuint32_t对应 Ascend C 原型__aicore__ inline uint64_t GetSize() const__aicore__ inline uint32_t GetSize() const初始化前提需先通过set_global_buffer传入全局数据地址通过 TQue / LocalMemAllocator 等方式获得值得注意的是二者返回的都是元素个数而非字节数若需字节数应结合元素个数与dtype的字节宽度自行换算或参考get_phy_addr、shape 信息等辅助接口。编译链路从 Python API 到 Ascend C 代码生成get_size的调用在 pyasc 中并不是简单的函数调用而是经过一层 IR中间表示的建图与代码生成。从源码结构看其完整链路如下Python 侧 API 定义在 python/asc/language/core/tensor.py 中get_size被require_jit与set_tensor_docstring(tensor_nameGlobalTensor, api_nameget_size)装饰表明它必须在 JIT 编译上下文中使用且其 docstring 由文档生成机制统一注入require_jit set_tensor_docstring(tensor_nameGlobalTensor, api_nameget_size) def get_size(self) - RuntimeInt: if self.shape is None: builder global_builder.get_ir_builder() handle builder.create_asc_GlobalTensorGetSizeOp(builder.get_ui64_type(), self.to_ir()) return PlainValue(handle) return math.prod(self.shape)IR 建图当GlobalTensor未显式设置 shape 时self.shape is None调用builder.create_asc_GlobalTensorGetSizeOp创建asc.global_tensor.get_size算子节点结果类型为ui64。这一 Op 在 TableGen 定义中声明为对 Ascend CGlobalTensor::GetSize方法的调用见 include/ascir/Dialect/Asc/IR/Core/Tensor.tddef AscendC_GlobalTensorGetSizeOp : APIOpglobal_tensor.get_size, GetSize, [AscMemberFunc] { let summary Call AscendC::GlobalTensor::GetSize method; let arguments (ins AscendC_GlobalTensor:$tensor); let results (outs UI64:$value); }代码生成随后由代码发射器Emit将 IR 转换为最终的内核 C 代码即文档中给出的__aicore__ inline uint64_t GetSize() const调用形式。一个值得注意的编译期优化在 python/asc/language/core/tensor.py 的实现中若GlobalTensor已经具备 shape 信息self.shape非空get_size会直接返回math.prod(self.shape)——即各维度之积而不再生成任何 IR 节点。这意味着 shape 已知时元素个数可以在编译期常量折叠降低运行时开销。这也解释了为何get_shape_info文档中强调Shape 信息没有默认值只有通过SetShapeInfo设置过 Shape 信息后才可以调用该接口获取正确的 ShapeInfo见 GlobalTensor.get_shape_info 文档。初始化前提set_global_buffer与元素个数的来源get_size返回的元素个数来源于初始化时通过set_global_buffer传入的buffer_size。set_global_buffer的完整签名如下见 GlobalTensor.set_global_buffer 文档GlobalTensor.set_global_buffer(buffer: GlobalAddress | None None) → None GlobalTensor.set_global_buffer(buffer: GlobalAddress | None None, buffer_size: int | None None) → NonebufferHost 侧传入的全局数据指针即GlobalAddress对应 Ascend C 的__gm__ PrimType* buffer。buffer_sizeGlobalTensor 所包含的类型为PrimType的数据个数需自行保证不会超出实际数据的长度对应 Ascend C 的uint64_t bufferSize。set_global_buffer对应的两个 Ascend C 重载原型为__aicore__ inline void SetGlobalBuffer(__gm__ PrimType* buffer, uint64_t bufferSize) __aicore__ inline void SetGlobalBuffer(__gm__ PrimType* buffer)在 Python 实现中python/asc/language/core/tensor.py若buffer为空或缺少数据类型会抛出ValueError传入buffer_size时生成带尺寸的GlobalTensorSetGlobalBufferOp否则生成不带尺寸的版本。从源码结构看get_size返回的元素个数正是基于此buffer_size维护的运行时信息。关键约束仅传入指针时元素个数为 0原文档明确给出的约束是使用仅传入全局数据指针的set_global_buffer接口对GlobalTensor进行初始化通过本接口获取到的元素个数为 0。即如果只调用set_global_buffer(x)不带buffer_size后续get_size()返回 0因为此时并没有记录元素个数的信息。因此凡是后续需要依赖元素个数如作为data_copy的拷贝长度、循环边界等的场景必须使用带buffer_size的set_global_buffer(x, size)形式完成初始化。实战示例在内核中获取并利用全局元素个数标准用法显式传入 buffer_size参考set_global_buffer文档中的调用示例以及 python/test/unit/language/core/test_global_tensor.py 中的单元测试标准写法如下import asc asc.jit def kernel_get_size(x: asc.GlobalAddress) - None: x_gm asc.GlobalTensor() x_gm.set_global_buffer(x, 8192) # 显式声明元素个数 x_size x_gm.get_size() # 返回 8192将get_size与数据搬运结合即可安全地驱动全局到局部的拷贝data_size 256 input_global asc.GlobalTensor() input_global.set_global_buffer(src_gm, data_size) input_local in_queue_x.alloc_tensor(asc.int32) asc.data_copy(input_local, input_global, data_size)此处input_global.get_size()与data_size语义一致元素个数可以作为data_copy长度的动态来源从而避免硬编码。完整的内核骨架结合 pyasc 的 JIT 框架一个可运行的内核骨架如下示意结构省略队列初始化细节import asc asc.jit def add_kernel(x: asc.GlobalAddress, y: asc.GlobalAddress, z: asc.GlobalAddress, size: asc.int32) - None: x_gm asc.GlobalTensor() x_gm.set_global_buffer(x, size) y_gm asc.GlobalTensor() y_gm.set_global_buffer(y, size) z_gm asc.GlobalTensor() z_gm.set_global_buffer(z, size) # 使用 get_size 获取元素个数驱动搬运与循环 total x_gm.get_size() # 通过 TQue 申请 LocalTensor 并执行数据拷贝、向量计算 # 具体队列/内存管理方式请参考 examples/ 目录下的完整算子示例更多可直接运行的端到端示例可参考仓库 examples 目录下的算子实现例如 01_add/add.py 演示了完整的asc.jit内核、Host 侧数据准备与 launch 流程。单元测试验证仓库中已有针对get_size的单元测试python/test/unit/language/core/test_global_tensor.pydef test_get_size(mock_launcher_run): asc.jit def kernel_get_size(x: asc.GlobalAddress) - None: x_gm asc.GlobalTensor() x_gm.set_global_buffer(x) x_size x_gm.get_size() data MockTensor(asc.float32) kernel_get_size1 assert mock_launcher_run.call_count 1该测试在config.set_platform(config.Backend.Model, checkFalse)的 Model 平台下运行使用MockTensor模拟输入验证内核可以正常编译、建图并 launch。注意此用例演示的是不带buffer_size的初始化路径对应元素个数为 0的约束场景用于验证接口可用性实际业务代码中应根据上述约束显式传入尺寸。常见误用与注意事项误以为返回字节数get_size返回元素个数。若需字节数需乘以dtype的字节宽度例如asc.float32为 4 字节。漏传buffer_size导致返回 0仅调用set_global_buffer(x)时get_size()返回 0无法用于数据拷贝长度或循环边界务必使用set_global_buffer(x, size)。buffer_size超出实际数据长度文档明确要求需自行保证不会超出实际数据的长度越界访问可能导致未定义行为需在 Host 侧确保尺寸正确。与get_shape_info混淆get_size返回元素个数标量get_shape_info返回ShapeInfo对象维度、原始 shape、data_format 等结构信息。且 shape 信息没有默认值必须先set_shape_info才能获得正确结果而get_size在 shape 已知时self.shape非空会直接做编译期常量折叠。返回值类型差异Python 侧统一表现为int但生成的内核代码中GlobalTensor为uint64_t、LocalTensor为uint32_t在大数据量场景下注意语义差异。相关接口与延伸阅读GlobalTensor的完整成员方法列表见 core.mdget_size与以下接口协同使用GlobalTensor.set_global_buffer初始化全局 Tensor 并声明元素个数get_size数据来源。GlobalTensor.get_phy_addr获取全局数据地址支持偏移量。GlobalTensor.get_shape_info获取维度/原始 shape/数据格式等结构信息。GlobalTensor.get_value按偏移读取指定位置的值。GlobalTensor.set_value按偏移写入指定位置的值。LocalTensor.get_sizeLocal Memory 侧的元素个数获取对照阅读可加深理解。上述接口对应的 IR Op 定义集中在 include/ascir/Dialect/Asc/IR/Core/Tensor.tdPython 实现集中在 python/asc/language/core/tensor.py单元测试见 python/test/unit/language/core/test_global_tensor.py可作为深入研读的起点。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表