lingvo
lingvo copied to clipboard
CUDA Support
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?
You are right that indeed there is no support for cuda. Briefly looking at the tensorflow code, it works like this:
- You specify
tf_kernel_library
with an appropriately named file<name>.cu.cc
that's picked up by a glob. - The file is build with
tf_kernel_library
. - That calls into
cuda_library
which is none other thannative.cc_library
withcuda_default_copts
. - There are similar things for the link version.
- bazel honors the
CC
andCXX
environment variables, so you might be able to use that to specifynvcc
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.