search

NVIDIA / nvmath-python

NVIDIA Math Libraries for the Python Ecosystem
https://docs.nvidia.com/cuda/nvmath-python/
Apache License 2.0
199 stars 8 forks source link

Support cuTensor? #3

Open wxj6000 opened 3 months ago

wxj6000 commented 3 months ago

I haven't see the complete documentation of this project. But the binding for cuTensor is not found in the repo. I wonder if there is any plan for supporting general tensor contractions?

leofang commented 3 months ago

Hi @wxj6000 thanks for reaching out and opening the very first issue in this repo! 😄

cuTENSOR support is on our roadmap but there is no ETA yet. I am sure you are aware that both CuPy and cuQuantum Python support cuTENSOR. We would be very curious to learn what use cases that you have in mind are not already covered by these projects. It would help us prioritize for new features.

FYI, the documentation will be live next week.

wxj6000 commented 3 months ago

@leofang Thanks. I was trying to understand the scope of the project and show my interest.

CuPy and cuQuantum Python do have the support for cuTensor. My collaborator and I am considering the kernel fusion of two tensor contraction kernels. For example, 

c = contract('ijk,ko->ijo', a, b)
d = contract('ijo,ilo->jl', c, c)