Ascend Extension for PyTorch的源码解析

本文涉及的产品
全局流量管理 GTM,标准版 1个月
公共DNS(含HTTPDNS解析),每月1000万次HTTP解析
云解析 DNS,旗舰版 1个月
简介: 本文介绍了Ascend对PyTorch代码的适配过程,包括源码下载、编译步骤及常见问题,详细解析了torch-npu编译后的文件结构和三种实现昇腾NPU算子调用的方式:通过torch的register方式、定义算子方式和API重定向映射方式。这对于开发者理解和使用Ascend平台上的PyTorch具有重要指导意义。

1 源码下载

Ascend对pytorch代码的适配,可从以下链接中获取。
Ascend/pytorch
执行如下命令即可。

git clone https://gitee.com/ascend/pytorch.git

2 目录结构解析

源码下载后,如果需要编译torch-npu,最好保持pytorch的源码版本匹配,以及其编译环境的gcc,g++等与torch-npu的版本匹配,否则会出现各种乱起八糟的问题。

执行编译命令:bash ci/build.sh --python=3.x

如:


csrc/aten/AutoCastOps.cpp:28:70: error: macro "KERNEL_PRIVATEUSEONE" passed 3 arguments, but takes just 2
KERNEL_PRIVATEUSEONE(_convolution, deprecated, lower_precision_fp)

在torch-npu编译成功之后,通过generate_code.sh会生成如下文件:

    torch_npu/csrc/aten/ADInplaceOrViewTypeEverything.cpp
    torch_npu/csrc/aten/ADInplaceOrViewType_0.cpp
    torch_npu/csrc/aten/ADInplaceOrViewType_1.cpp
    torch_npu/csrc/aten/CustomFunctions.cpp
    torch_npu/csrc/aten/CustomFunctions.h
    torch_npu/csrc/aten/CustomRedispatch.cpp
    torch_npu/csrc/aten/CustomRedispatch.h
    torch_npu/csrc/aten/CustomRegisterSchema.cpp
    torch_npu/csrc/aten/ForeachRegister.cpp
    torch_npu/csrc/aten/Functions.cpp
    torch_npu/csrc/aten/Functions.h
    torch_npu/csrc/aten/NPUOpApiNativeFunctions.h
    torch_npu/csrc/aten/QuantizedRegister.cpp
    torch_npu/csrc/aten/RegisterFunctionalizationEverything.cpp
    torch_npu/csrc/aten/RegisterFunctionalization_0.cpp
    torch_npu/csrc/aten/RegisterFunctionalization_1.cpp
    torch_npu/csrc/aten/RegisterSparseCsrNPU.cpp
    torch_npu/csrc/aten/RegisterSparseNPU.cpp
    torch_npu/csrc/aten/VariableType.h
    torch_npu/csrc/aten/VariableTypeEverything.cpp
    torch_npu/csrc/aten/VariableType_0.cpp
    torch_npu/csrc/aten/npu_native_functions_by_codegen.yaml
    torch_npu/csrc/aten/python_functions.h
    torch_npu/csrc/aten/python_functionsEverything.cpp
    torch_npu/csrc/aten/python_functions_0.cpp
    torch_npu/csrc/aten/python_functions_1.cpp
    torch_npu/csrc/aten/variable_factories.h
    torch_npu/testing/_npu_testing_utils.py
    torch_npu/utils/custom_ops.py
    torch_npu/utils/exposed_api.py

上述文件生成路径默认的是torch_npu/csrc/aten。算子编译信息的yaml文件:torch_npu/csrc/aten/npu_native_functions.yaml

打开上述的的文件中,从中分析可知大概有3种方式实现昇腾npu算子的调用。

3. 算子注册方式

本质上,ascend上对pytroch框架的适配代码,主要是将npu上的算子库对接起来。如何对接这些算子,是一套机制的问题,本身应该不复杂。

3.1 通过torch的regsiter方式

直接调用npu的算子。torch_npu/csrc/aten/RegisterSparseNPU.cpp

TORCH_LIBRARY_IMPL(aten, SparsePrivateUse1, m) {
   
m.impl("abs", TORCH_FN(wrap_SparseNPU_abs_));
m.impl("abs_", TORCH_FN(wrap_SparseNPU_abs__));
m.impl("abs.out", TORCH_FN(wrap_SparseNPU_abs_out));
m.impl("sgn", TORCH_FN(wrap_SparseNPU_sgn_));
m.impl("sgn_", TORCH_FN(wrap_SparseNPU_sgn__));
m.impl("sgn.out", TORCH_FN(wrap_SparseNPU_sgn_out));

3.2 通过定义算子方式

参考文件:torch_npu/csrc/aten/CustomFunctions.cpp

#include <ATen/core/dispatch/Dispatcher.h>

#include "torch_npu/csrc/aten/CustomFunctions.h"


namespace at_npu {
   
namespace native {
   
namespace custom_ops {
   

int64_t npu_change_data_ptr(const at::Tensor & dst, const at::Tensor & src, int64_t index) {
   
    static auto op = c10::Dispatcher::singleton().findSchemaOrThrow("npu::npu_change_data_ptr", "").typed<int64_t (const at::Tensor &, const at::Tensor &, int64_t)>();
    return op.call(dst, src, index);
}
int64_t get_npu_format(const at::Tensor & self) {
   
    static auto op = c10::Dispatcher::singleton().findSchemaOrThrow("npu::get_npu_format", "").typed<int64_t (const at::Tensor &)>();
    return op.call(self);
}
at::Tensor npu_format_cast(const at::Tensor & self, const at::Tensor & dst) {
   
    static auto op = c10::Dispatcher::singleton().findSchemaOrThrow("npu::npu_format_cast", "Tensor").typed<at::Tensor (const at::Tensor &, const at::Tensor &)>();
    return op.call(self, dst);
}
at::Tensor & npu_format_cast_(at::Tensor & self, int64_t acl_format) {
   
    static auto op = c10::Dispatcher::singleton().findSchemaOrThrow("npu::npu_format_cast_", "acl_format").typed<at::Tensor & (at::Tensor &, int64_t)>();
    return op.call(self, acl_format);

 at::Tensor & npu_format_cast_(at::Tensor & self, const at::Tensor & src) {
   
    static auto op = c10::Dispatcher::singleton().findSchemaOrThrow("npu::npu_format_cast_", "").typed<at::Tensor & (at::Tensor &, const at::Tensor &)>();
    return op.call(self, src);
}
at::Tensor empty_with_format(at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, int64_t acl_format) {
   
    static auto op = c10::Dispatcher::singleton().findSchemaOrThrow("npu::empty_with_format", "").typed<at::Tensor (at::IntArrayRef, ::std::optional<at::ScalarType>, ::std::optional<at::Layout>, ::std::optional<at::Device>, ::std::optional<bool>, int64_t)>();
    return op.call(size, dtype, layout, device, pin_memory, acl_format);
}
at::Tensor unsafe_empty_with_format(at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, int64_t acl_format, bool keep_format) {
   
    static auto op = c10::Dispatcher::singleton().findSchemaOrThrow("npu::unsafe_empty_with_format", "").typed<at::Tensor (at::IntArrayRef, ::std::optional<at::ScalarType>, ::std::optional<at::Layout>, ::std::optional<at::Device>, ::std::optional<bool>, int64_t, bool)>();
    return op.call(size, dtype, layout, device, pin_memory, acl_format, keep_format);
}
 ~/pytorch-ascend/torch_npu/csrc/aten/CustomFunctions.cpp[1,RO]  

...

}
}
}

3.3 通过API重定向映射的方式

参考文件:torch_npu/utils/custom_ops.py

torch_npu.npu_layer_norm_eval = torch.ops.npu.npu_layer_norm_eval
torch_npu.npu_fused_attention_score_grad = torch.ops.npu.npu_fused_attention_score_grad
torch_npu.npu_quant_conv2d = torch.ops.npu.npu_quant_conv2d
torch_npu.npu_view_copy = torch.ops.npu.npu_view_copy
torch_npu.npu_fast_gelu = torch.ops.npu.npu_fast_gelu
torch_npu.npu_fused_attention_layernorm_qkv_fwd = torch.ops.npu.npu_fused_attention_layernorm_qkv_fwd
torch_npu.npu_fast_gelu_backward = torch.ops.npu.npu_fast_gelu_backward
torch_npu.npu_bmm_v2_mat1_backward = torch.ops.npu.npu_bmm_v2_mat1_backward

以上属于个人理解,如有错误敬请指正。

相关文章
|
1月前
|
监控 Java 应用服务中间件
高级java面试---spring.factories文件的解析源码API机制
【11月更文挑战第20天】Spring Boot是一个用于快速构建基于Spring框架的应用程序的开源框架。它通过自动配置、起步依赖和内嵌服务器等特性,极大地简化了Spring应用的开发和部署过程。本文将深入探讨Spring Boot的背景历史、业务场景、功能点以及底层原理,并通过Java代码手写模拟Spring Boot的启动过程,特别是spring.factories文件的解析源码API机制。
71 2
|
17天前
|
人工智能 PyTorch 算法框架/工具
Ascend Extension for PyTorch是个what?
Ascend Extension for PyTorch 是针对华为昇腾处理器的PyTorch框架适配插件,旨在让PyTorch开发者能充分利用昇腾AI处理器的强大计算能力。此扩展通过最小化对原生PyTorch的改动,实现了对昇腾NPU的支持,包括动态图特性、自动微分等功能的完整继承,并提供了与原生PyTorch一致的使用体验。项目详情及源码可在昇腾社区获取。
Ascend Extension for PyTorch是个what?
|
19天前
|
机器学习/深度学习 人工智能 PyTorch
Transformer模型变长序列优化:解析PyTorch上的FlashAttention2与xFormers
本文探讨了Transformer模型中变长输入序列的优化策略,旨在解决深度学习中常见的计算效率问题。文章首先介绍了批处理变长输入的技术挑战,特别是填充方法导致的资源浪费。随后,提出了多种优化技术,包括动态填充、PyTorch NestedTensors、FlashAttention2和XFormers的memory_efficient_attention。这些技术通过减少冗余计算、优化内存管理和改进计算模式,显著提升了模型的性能。实验结果显示,使用FlashAttention2和无填充策略的组合可以将步骤时间减少至323毫秒,相比未优化版本提升了约2.5倍。
35 3
Transformer模型变长序列优化:解析PyTorch上的FlashAttention2与xFormers
|
22天前
|
缓存 监控 Java
Java线程池提交任务流程底层源码与源码解析
【11月更文挑战第30天】嘿,各位技术爱好者们,今天咱们来聊聊Java线程池提交任务的底层源码与源码解析。作为一个资深的Java开发者,我相信你一定对线程池并不陌生。线程池作为并发编程中的一大利器,其重要性不言而喻。今天,我将以对话的方式,带你一步步深入线程池的奥秘,从概述到功能点,再到背景和业务点,最后到底层原理和示例,让你对线程池有一个全新的认识。
50 12
|
1月前
|
存储 安全 Linux
Golang的GMP调度模型与源码解析
【11月更文挑战第11天】GMP 调度模型是 Go 语言运行时系统的核心部分,用于高效管理和调度大量协程(goroutine)。它通过少量的操作系统线程(M)和逻辑处理器(P)来调度大量的轻量级协程(G),从而实现高性能的并发处理。GMP 模型通过本地队列和全局队列来减少锁竞争,提高调度效率。在 Go 源码中,`runtime.h` 文件定义了关键数据结构,`schedule()` 和 `findrunnable()` 函数实现了核心调度逻辑。通过深入研究 GMP 模型,可以更好地理解 Go 语言的并发机制。
|
1月前
|
消息中间件 缓存 安全
Future与FutureTask源码解析,接口阻塞问题及解决方案
【11月更文挑战第5天】在Java开发中,多线程编程是提高系统并发性能和资源利用率的重要手段。然而,多线程编程也带来了诸如线程安全、死锁、接口阻塞等一系列复杂问题。本文将深度剖析多线程优化技巧、Future与FutureTask的源码、接口阻塞问题及解决方案,并通过具体业务场景和Java代码示例进行实战演示。
54 3
|
2月前
|
存储
让星星⭐月亮告诉你,HashMap的put方法源码解析及其中两种会触发扩容的场景(足够详尽,有问题欢迎指正~)
`HashMap`的`put`方法通过调用`putVal`实现,主要涉及两个场景下的扩容操作:1. 初始化时,链表数组的初始容量设为16,阈值设为12;2. 当存储的元素个数超过阈值时,链表数组的容量和阈值均翻倍。`putVal`方法处理键值对的插入,包括链表和红黑树的转换,确保高效的数据存取。
63 5
|
2月前
|
Java Spring
Spring底层架构源码解析(三)
Spring底层架构源码解析(三)
148 5
|
2月前
|
算法 PyTorch 算法框架/工具
Pytorch学习笔记(九):Pytorch模型的FLOPs、模型参数量等信息输出(torchstat、thop、ptflops、torchsummary)
本文介绍了如何使用torchstat、thop、ptflops和torchsummary等工具来计算Pytorch模型的FLOPs、模型参数量等信息。
363 2
|
2月前
|
机器学习/深度学习 自然语言处理 监控
利用 PyTorch Lightning 搭建一个文本分类模型
利用 PyTorch Lightning 搭建一个文本分类模型
69 8
利用 PyTorch Lightning 搭建一个文本分类模型
下一篇
DataWorks