pytorch 是如何调用 cusolver API 的调用

0,环境

ubuntu 22.04

pytorch 2.3.1

x86

RTX 3080

cuda 12.2

1, 示例代码

以potrs为例;

hello_cholesk.py

复制代码
""" 
hello_cholesky.py
step1, Cholesky decompose;
step2, inverse A;
step3, Cholesky again;
python3 hello_cholesky.py --size 256  --cuda_device_id  0
"""
import torch
import time
import argparse


def cholesky_measure(A, cuda_dev=0):
    dev = torch.device(f"cuda:{cuda_dev}")
    A = A.to(dev)

    print(f'Which device to compute : {dev}')
  
    SY = 100* torch.mm(A, A.t()) +  200*torch.eye(N, device=dev)

    to_start = time.time() 
    SY = torch.linalg.cholesky(SY)
    SY = torch.cholesky_inverse(SY)
    SY = torch.linalg.cholesky(SY, upper=True)
    run_time = time.time() - to_start   
     
    print(f'The device: {dev}, run: {run_time:.3f} second')
    print(f'SY : {SY}')
    print(f'****'*20)

    return run_time

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='dim of A.')
    parser.add_argument('--N', type=int, default=512, required=True, help='dim of A')
    args = parser.parse_args()
    N = args.N

    print(f'A N : {N}')    
    A = torch.randn(N, N)
       
    cuda_dev = 0
    time_dev0 = cholesky_measure(A, cuda_dev)    
    time_dev1 = cholesky_measure(A, cuda_dev+1)    
    print(f'time_dev0 /time_dev1 = {time_dev0/time_dev1:.2f} ')

运行效果:

2,调用栈跟踪

跟踪如下调用关系:

复制代码
Tensor cholesky_inverse(const Tensor &input, bool upper)    aten/src/ATen/native/BatchLinearAlgebra.cpp
	static Tensor& cholesky_inverse_out_info(Tensor& result, Tensor& infos, const Tensor& input, bool upper)
	DECLARE_DISPATCH(cholesky_inverse_fn, cholesky_inverse_stub);
	REGISTER_ARCH_DISPATCH(cholesky_inverse_stub, DEFAULT, &cholesky_inverse_kernel_impl);
	Tensor& cholesky_inverse_kernel_impl(Tensor &result, Tensor& infos, bool upper)
	Tensor& cholesky_inverse_kernel_impl_cusolver(Tensor &result, Tensor& infos, bool upper)
	void _cholesky_inverse_cusolver_potrs_based(Tensor& result, Tensor& infos, bool upper)
	template<typename scalar_t>
	inline static void apply_cholesky_cusolver_potrs(Tensor& self_working_copy, const Tensor& A_column_major_copy, bool upper, Tensor& infos)
	at::cuda::solver::potrs<scalar_t>(
      handle, uplo, n_32, nrhs_32,
      A_ptr + i * A_matrix_stride,
      lda_32,
      self_working_copy_ptr + i * self_matrix_stride,
      ldb_32,
      infos_ptr
    );

一些细节:

相关推荐
后端小肥肠3 分钟前
躺赚必备!RPA+Coze+豆包:公众号自动发文,AI率0%亲测有效(附AI率0%提示词)
人工智能·aigc·coze
摘星编程14 分钟前
CloudBase AI ToolKit实战:从0到1开发一个智能医疗网站
人工智能·腾讯云·ai代码远征季#h5应用·ai医疗应用·cloudbase开发
锅挤19 分钟前
深度学习5(深层神经网络 + 参数和超参数)
人工智能·深度学习·神经网络
一支烟一朵花23 分钟前
630,百度文心大模型4.5系列开源!真香
人工智能·百度·开源·文心一言
网安INF25 分钟前
深层神经网络:原理与传播机制详解
人工智能·深度学习·神经网络·机器学习
喜欢吃豆30 分钟前
目前最火的agent方向-A2A快速实战构建(二): AutoGen模型集成指南:从OpenAI到本地部署的全场景LLM解决方案
后端·python·深度学习·flask·大模型
AIbase20241 小时前
国内MCP服务平台推荐!aibase.cn上线MCP服务器集合平台
运维·服务器·人工智能
喜欢吃豆2 小时前
快速手搓一个MCP服务指南(九): FastMCP 服务器组合技术:构建模块化AI应用的终极方案
服务器·人工智能·python·深度学习·大模型·github·fastmcp
星融元asterfusion2 小时前
基于路径质量的AI负载均衡异常路径检测与恢复策略
人工智能·负载均衡·异常路径