如何在Tensorflow BUILD文件中添加nvcc的-rdc = true选项以构建我的自定义操作?

时间:2018-10-27 15:32:31

标签: tensorflow

在我定制的基于gpu的op中,有必要实现NVIDIA Dynamic Parallelism机制。也就是说,我要实现内核函数的嵌套调用。我知道要实现上述机制,应将'-rdc = true'选项传递给nvcc编译器。

nvcc -o a nestedKernelFun.cu.cc -arch=sm_35 -rdc=true -lcudadev

但是,如何将'-rdc = true'添加到自定义op BUILD文件中?当我直接将副本添加到tf_custom_op_library中时:

tf_custom_op_library(
name = "my_op.so",
srcs = ["nestedKernelFun.cu.cc"],
copts = [-rdc=true],
)

bazel提示我tf_custom_op_library不支持复制。

非常感谢。

0 个答案:

没有答案