Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[FEA] Allow cublas wrapper to work on different input/output dtypes #977

Open
benfred opened this issue Nov 1, 2022 · 0 comments
Open
Labels
feature request New feature or request

Comments

@benfred
Copy link
Member

benfred commented Nov 1, 2022

Is your feature request related to a problem? Please describe.
Right now we are constrained to using the same input and output dtypes for all inputs and outputs in the cublas wrapper.
Wrapping the cublas*Ex api's would allow us to partially overcome the restriction and work with mixed dtypes between input and output.

As an example, we're wrapping the cublasSgemm function - which requires all inputs to be a float32, whereas if we changed to wrapping the cublasGemmEx function we could avoid this limitation and have the inputs be a float16 and the output be float32.

Describe the solution you'd like
In detail/cublas_wrapper.cuh we could change:

  • cublas<t>gemm to cublasGemmEx
  • cublas<T>gemmBatched to cublasgemmBatchedEx
  • cublas<t>axpy to cublasaxpyEx
  • cublas<t>dot to cublasDotEx
    etc

setting the cudaDataType parameters in the cublas*Ex calls automatically from the template parameters.

Additional context
One issue is that the cublasEx api's don't support all combinations of type parameters. As an example, the cublasGemmEx api requires the A/B input parameters to have the same type afaict - despite taking two different datatype parameters for these values. https://docs.nvidia.com/cuda/cublas/index.html#cublas-GemmEx

Likewise the supported datatype combinations for cublasDotEx requires both inputs and output to have the same type - meaning this change won't have much benefit currently for the linalg::dot function. https://docs.nvidia.com/cuda/cublas/index.html#cublas-dotEx

@benfred benfred added the feature request New feature or request label Nov 1, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
feature request New feature or request
Projects
None yet
Development

No branches or pull requests

1 participant