lingvo icon indicating copy to clipboard operation
lingvo copied to clipboard

CUDA Support

Open durson opened this issue 4 years ago • 1 comments

In order to add custom op implementation I need CUDA support. However I see no possibility to compile .cu files in Lingvo. Bazel has no CUDA support AFAIK, so some backward shenanigans are necessary. Have anyone here made this work or can point me to a solution?

durson avatar May 05 '20 09:05 durson

You are right that indeed there is no support for cuda. Briefly looking at the tensorflow code, it works like this:

  1. You specify tf_kernel_library with an appropriately named file <name>.cu.cc that's picked up by a glob.
  2. The file is build with tf_kernel_library.
  3. That calls into cuda_library which is none other than native.cc_library with cuda_default_copts.
  4. There are similar things for the link version.
  5. bazel honors the CC and CXX environment variables, so you might be able to use that to specify nvcc as your compiler.

To get the true value of commands in tensorflow, if you have a source version lying around, it's easy to start with a working version, than make a file fail and see the command-line with --verbose_failures.

In lingvo, we have opted to rebuild the environment with the same name as tensorflow, prefixing with lingvo_, for instance, tf_cc_library becomes lingvo_cc_library. That's all defined in lingvo.bzl. We try to keep things simple and minimal.

So, I'd start with setting the CXX to nvcc and building a lingvo_cc_library specifying the additional copts by hand, and adding the linkopts by hand as well. If that works, then we are all good.

It's likely to fail and we might have to import cuda_configure.bzl in one way or another.

Another path is to build a custom op and loaded from your python code.

drpngx avatar May 05 '20 15:05 drpngx