how to use Composable Kernel library in tensorflow?

[south-ocean]

Issue type

Feature Request

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

tf

Custom code

Yes

OS platform and distribution

No response

Mobile device

No response

Python version

No response

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

I see that AMD has a Composable Kernel library, and I want to ask if tensorflow already supports this library? If not, how can you modify it to make TensorFlow call the Composable Kernel library? Are there any relevant introduction and usage documents?

Standalone code to reproduce the issue

As mentioned above

Relevant log output

As mentioned abov

[wenchenvincent]

Hi, currently, Tensorflow ROCm does not call Composable Kernel directly. In the future, it is possible that we also use CK as upstream Tensorflow might use CUTLASS. However, the user of tensorflow does not usually need to worry about it as it is all under the hood.

[south-ocean]

Thanks, as you said, my understanding is that CK will be connected to TensorFlow through the library like cutlass? This allows miopen and blas to offload part of the CK.

[wenchenvincent]

Thanks, as you said, my understanding is that CK will be connected to TensorFlow through the library like cutlass? This allows miopen and blas to offload part of the CK.

Actually CK is currently used as one of the backends for MIOpen already. And Tensorflow uses MIOpen. In the future, we might add connectors that would use CK directly from Tensorflow.

[south-ocean]

thanks