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 分钟前
基于多摄像头融合的智能小车自动驾驶系统完整实现
人工智能·机器学习·自动驾驶
rengang6622 分钟前
351-Spring AI Alibaba Dashscope 多模型示例
java·人工智能·spring·多模态·spring ai·ai应用编程
NewCarRen24 分钟前
整合STPA、ISO 26262与SOTIF的自动驾驶安全需求推导与验证
人工智能·安全·自动驾驶·预期功能安全
嵌入式老牛36 分钟前
1.0 机器学习实际上是什么
人工智能·机器学习
易晨 微盛·企微管家37 分钟前
汽车行业SCRM:企业微信+服务商模式破解汽车服务行业痛点的案例分析
大数据·人工智能·汽车·产品运营·企业微信
易晨 微盛·企微管家42 分钟前
生鲜电商企业微信私域代运营:从去中心化运营看微盛AI·企微管家SCRM适配案例
人工智能·微信·企业微信
李元豪1 小时前
MetaGPT、AutoGen、XAgent camel仔细对比
人工智能
AndrewHZ1 小时前
【图像处理基石】图像Inpainting入门详解
图像处理·人工智能·深度学习·opencv·transformer·图像修复·inpainting
jiushun_suanli1 小时前
PyTorch CV模型实战全流程(一)
人工智能·pytorch·python
学技术的大胜嗷1 小时前
如何裁剪YOLOv8m的大目标检测头并验证其结构
深度学习·yolo·目标检测·计算机视觉