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
    );

一些细节:

相关推荐
华新嘉华DTC创新营销1 小时前
华新嘉华:AI搜索优化重塑本地生活行业:智能推荐正取代“关键词匹配”
人工智能·百度·生活
SmartBrain2 小时前
DeerFlow 实践:华为IPD流程的评审智能体设计
人工智能·语言模型·架构
l1t3 小时前
利用DeepSeek实现服务器客户端模式的DuckDB原型
服务器·c语言·数据库·人工智能·postgresql·协议·duckdb
寒月霜华4 小时前
机器学习-数据标注
人工智能·机器学习
九章云极AladdinEdu5 小时前
超参数自动化调优指南:Optuna vs. Ray Tune 对比评测
运维·人工智能·深度学习·ai·自动化·gpu算力
人工智能训练师6 小时前
Ubuntu22.04如何安装新版本的Node.js和npm
linux·运维·前端·人工智能·ubuntu·npm·node.js
cxr8288 小时前
SPARC方法论在Claude Code基于规则驱动开发中的应用
人工智能·驱动开发·claude·智能体
研梦非凡8 小时前
ICCV 2025|从粗到细:用于高效3D高斯溅射的可学习离散小波变换
人工智能·深度学习·学习·3d
幂简集成8 小时前
Realtime API 语音代理端到端接入全流程教程(含 Demo,延迟 280ms)
人工智能·个人开发
龙腾-虎跃9 小时前
FreeSWITCH FunASR语音识别模块
人工智能·语音识别·xcode